code: def train_batch(
      self, batch: pytorch.TorchData, epoch_idx: int, batch_idx: int
  ) -> Dict[str, torch.Tensor]:
      data, label = batch
  
      loss = self.loss_fn(self.model(data), label)
  
      self.context.backward(loss)
  
      self.context.step_optimizer(self.opt)
  
      return {
          "loss": loss,
      }
  
file_path: determined/e2e_tests/tests/fixtures/pytorch_identity/model_def.py
imports: from typing import Any, Dict, Tuple
  import torch.utils.data
  from determined import pytorch
label_desc: Use image node_id 2 for calling the IdentityPyTorchTrial obj's underlying member method code with example usage: obj.train_batch(batch, epoch_idx, batch_idx) and returns: dict