def create_trainer(model, optimizer, criterion, lr_scheduler, train_sampler, config, logger):
    device = idist.device()

    def train_step(engine, batch):
        (x, y) = (batch[0], batch[1])
        if (x.device != device):
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
        model.train()
        y_pred = model(x)
        loss = criterion(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if ((config['log_every_iters'] > 0) and (((engine.state.iteration - 1) % config['log_every_iters']) == 0)):
            batch_loss = loss.item()
            engine.state.saved_batch_loss = batch_loss
        else:
            batch_loss = engine.state.saved_batch_loss
        return {'batch loss': batch_loss}
    trainer = Engine(train_step)
    trainer.state.saved_batch_loss = (- 1.0)
    trainer.state_dict_user_keys.append('saved_batch_loss')
    trainer.logger = logger
    to_save = {'trainer': trainer, 'model': model, 'optimizer': optimizer, 'lr_scheduler': lr_scheduler}
    metric_names = ['batch loss']
    common.setup_common_training_handlers(trainer=trainer, train_sampler=train_sampler, to_save=to_save, save_every_iters=config['checkpoint_every'], save_handler=get_save_handler(config), lr_scheduler=lr_scheduler, output_names=(metric_names if (config['log_every_iters'] > 0) else None), with_pbars=False, clear_cuda_cache=False)
    resume_from = config['resume_from']
    if (resume_from is not None):
        checkpoint_fp = Path(resume_from)
        assert checkpoint_fp.exists(), "Checkpoint '{}' is not found".format(checkpoint_fp.as_posix())
        logger.info('Resume from a checkpoint: {}'.format(checkpoint_fp.as_posix()))
        checkpoint = torch.load(checkpoint_fp.as_posix(), map_location='cpu')
        Checkpoint.load_objects(to_load=to_save, checkpoint=checkpoint)
    return trainer
def create_trainer(model, optimizer, criterion, lr_scheduler, train_sampler, config, logger):
    device = idist.device()

    def train_step(engine, batch):
        (x, y) = (batch[0], batch[1])
        if (x.device != device):
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
        model.train()
        y_pred = model(x)
        loss = criterion(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if ((config['log_every_iters'] > 0) and (((engine.state.iteration - 1) % config['log_every_iters']) == 0)):
            batch_loss = loss.item()
            engine.state.saved_batch_loss = batch_loss
        else:
            batch_loss = engine.state.saved_batch_loss
        return {'batch loss': batch_loss}
    trainer = Engine(train_step)
    trainer.state.saved_batch_loss = (- 1.0)
    trainer.state_dict_user_keys.append('saved_batch_loss')
    trainer.logger = logger
    to_save = {'trainer': trainer, 'model': model, 'optimizer': optimizer, 'lr_scheduler': lr_scheduler}
    metric_names = ['batch loss']
    common.setup_common_training_handlers(trainer=trainer, train_sampler=train_sampler, to_save=to_save, save_every_iters=config['checkpoint_every'], save_handler=get_save_handler(config), lr_scheduler=lr_scheduler, output_names=(metric_names if (config['log_every_iters'] > 0) else None), with_pbars=False, clear_cuda_cache=False)
    resume_from = config['resume_from']
    if (resume_from is not None):
        checkpoint_fp = Path(resume_from)
        assert checkpoint_fp.exists(), "Checkpoint '{}' is not found".format(checkpoint_fp.as_posix())
        logger.info('Resume from a checkpoint: {}'.format(checkpoint_fp.as_posix()))
        checkpoint = torch.load(checkpoint_fp.as_posix(), map_location='cpu')
        Checkpoint.load_objects(to_load=to_save, checkpoint=checkpoint)
    return trainer