def METHOD_NAME(self, iter_dataloader, epoch_idx, batch_idx) -> Dict[str, torch.Tensor]:
    batch = self.context.to_device(next(iter_dataloader))
    inputs, labels = batch[0], batch[1]
    if self.fp16:
        inputs = inputs.half()
    outputs = self.model_engine(inputs)
    loss = self.criterion(outputs, labels)
    self.model_engine.backward(loss)
    self.model_engine.step()
    return {"loss": loss.item()}