code: def train_batch(
      self, batch: TorchData, epoch_idx: int, batch_idx: int
  ) -> Dict[str, torch.Tensor]:
      batch = cast(Tuple[torch.Tensor, torch.Tensor], batch)
      data, labels = batch
  
      with autocast():
          output = self.model(data)
          loss = torch.nn.functional.nll_loss(output, labels)
  
      self.context.backward(self.scaler.scale(loss))
      self.context.step_optimizer(self.optimizer, scaler=self.scaler)
      self.scaler.update()
  
      return {"loss": loss}
  
file_path: determined/harness/tests/experiment/fixtures/pytorch_amp/manual_amp_model_def.py
imports: from typing import Any, Dict, Optional, Sequence, Tuple, Union, cast
  import torch
  from torch.cuda.amp import GradScaler, autocast
  from train import MNistTrial
  from determined.pytorch import PyTorchTrialContext
label_desc: Use image node_id 2 for calling the MNistManualAMPTrial obj's underlying member method code with example usage: obj.train_batch(batch, epoch_idx, batch_idx) and returns: dict