code: def __init__(self, context: pytorch.PyTorchTrialContext) -> None:
      self.context = context
  
      model = torch.nn.Linear(1, 1, False)
      model.weight.data.fill_(0)
      self.model = context.wrap_model(model)
  
      self.lr = 0.001
  
      optimizer = torch.optim.SGD(self.model.parameters(), self.lr)
      self.opt = context.wrap_optimizer(optimizer)
  
      self.loss_fn = torch.nn.MSELoss(reduction="mean")
      self.metrics_callback = MetricsCallback()
  
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 1 to create a new IdentityPyTorchTrial object from inherited base classes: pytorch with example: obj = IdentityPyTorchTrial(context)