code: def train_batch(
      self, batch: pytorch.TorchData, epoch_idx: int, batch_idx: int
  ) -> Dict[str, torch.Tensor]:
      batch_data, labels = batch
  
      output = self.model(batch_data)
      loss = self.loss_fn(output, labels)
  
      self.context.backward(loss)
      self.context.step_optimizer(self.optimizer)
  
      return {"loss": loss}
  
file_path: determined/examples/tutorials/mnist_pytorch/train.py
imports: import logging
  import pathlib
  from typing import Any, Dict
  import data
  import model
  import torch
  from ruamel import yaml
  from torch import nn
  import determined
  from determined import pytorch
label_desc: Use image node_id 4 for calling the MNistTrial obj's underlying member method code with example usage: obj.train_batch(batch, epoch_idx, batch_idx) and returns: dict