def train_step(self):
    'Train step to be executed in train loop.'
    self.last_batch = None
    with Timer() as timer:
        sample = next(self.data_iterator)
    assert (isinstance(sample, dict) and ('input' in sample) and ('target' in sample)), (f"Returned sample [{sample}] is not a map with 'input' and" + "'target' keys")
    target = sample['target']
    if self.use_gpu:
        sample = recursive_copy_to_gpu(sample, non_blocking=True)
    if (self.mixup_transform is not None):
        sample = self.mixup_transform(sample)
    torch_amp_context = (torch.cuda.amp.autocast() if (self.amp_type == AmpType.PYTORCH) else contextlib.suppress())
    do_step = self._should_do_step()
    ctx_mgr_model = (self.distributed_model.no_sync() if ((self.distributed_model is not None) and (not do_step)) else contextlib.suppress())
    ctx_mgr_loss = (self.distributed_loss.no_sync() if ((self.distributed_loss is not None) and (not do_step)) else contextlib.suppress())
    with ctx_mgr_model, ctx_mgr_loss:
        with torch.enable_grad(), torch_amp_context:
            output = self.model(sample['input'])
            local_loss = self.compute_loss(output, sample)
            loss = local_loss.detach().clone()
            self.losses.append((loss.data.cpu().item() * target.size(0)))
            self.update_meters(output, sample)
        self.run_optimizer(local_loss)
    self.num_updates += self.get_global_batchsize()
    self.last_batch = LastBatchInfo(loss=loss, output=output, target=target, sample=sample, step_data={'sample_fetch_time': timer.elapsed_time})
Train step to be executed in train loop.
def train_step(self):
    
    self.last_batch = None
    with Timer() as timer:
        sample = next(self.data_iterator)
    assert (isinstance(sample, dict) and ('input' in sample) and ('target' in sample)), (f"Returned sample [{sample}] is not a map with 'input' and" + "'target' keys")
    target = sample['target']
    if self.use_gpu:
        sample = recursive_copy_to_gpu(sample, non_blocking=True)
    if (self.mixup_transform is not None):
        sample = self.mixup_transform(sample)
    torch_amp_context = (torch.cuda.amp.autocast() if (self.amp_type == AmpType.PYTORCH) else contextlib.suppress())
    do_step = self._should_do_step()
    ctx_mgr_model = (self.distributed_model.no_sync() if ((self.distributed_model is not None) and (not do_step)) else contextlib.suppress())
    ctx_mgr_loss = (self.distributed_loss.no_sync() if ((self.distributed_loss is not None) and (not do_step)) else contextlib.suppress())
    with ctx_mgr_model, ctx_mgr_loss:
        with torch.enable_grad(), torch_amp_context:
            output = self.model(sample['input'])
            local_loss = self.compute_loss(output, sample)
            loss = local_loss.detach().clone()
            self.losses.append((loss.data.cpu().item() * target.size(0)))
            self.update_meters(output, sample)
        self.run_optimizer(local_loss)
    self.num_updates += self.get_global_batchsize()
    self.last_batch = LastBatchInfo(loss=loss, output=output, target=target, sample=sample, step_data={'sample_fetch_time': timer.elapsed_time})