code: def train_batch(self, batch, epoch_idx, batch_idx):
      loss = torch.nn.MSELoss()(self.model(batch), batch)
      self.context.step_optimizer(self.opt)
      return {"loss": loss}
  
file_path: determined/harness/tests/experiment/fixtures/ancient-checkpoints/0.13.13-pytorch-flex/code/model_def.py
imports: import torch
  from torch import nn
  from determined import pytorch
label_desc: Use image node_id 2 for calling the OneVarPytorchTrial obj's underlying member method code with example usage: obj.train_batch(batch, epoch_idx, batch_idx) and returns: dict