code: def __init__(
      self, context: pytorch.PyTorchTrialContext, hparams: Dict
  ) -> None:
      self.context = context
  
      # Trial-level constants.
      self.data_dir = pathlib.Path("data")
      self.data_url = "https://s3-us-west-2.amazonaws.com/determined-ai-test-data/pytorch_mnist.tar.gz"
  
      self.batch_size = 64
      self.per_slot_batch_size = (
          self.batch_size // self.context.distributed.get_size()
      )
  
      # Define loss function.
      self.loss_fn = nn.NLLLoss()
  
      # Define model.
      self.model = self.context.wrap_model(
          model.build_model(hparams=hparams)
      )
  
      # Configure optimizer.
      self.optimizer = self.context.wrap_optimizer(
          torch.optim.Adadelta(
              self.model.parameters(), lr=hparams["learning_rate"]
          )
      )
  
file_path: determined/harness/tests/experiment/fixtures/pytorch_amp/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 1 to create a new MNistTrial object from inherited base classes: pytorch with example: obj = MNistTrial(context, hparams)