code: def __init__(self, hidden_dim, frozen_weights):
      super(MyModel, self).__init__()
      self.act = torch.nn.ReLU()
      self.cel = torch.nn.CrossEntropyLoss()
      self.linears = torch.nn.ModuleList(
          [
              torch.nn.Linear(hidden_dim, 1),
              torch.nn.Linear(1, 1),
              torch.nn.Linear(1, hidden_dim),
          ]
      )
      if frozen_weights:
          self.linears[0].weight.requires_grad = False
          self.linears[0].bias.requires_grad = False
  
imports: import pytest
  import deepspeed.comm
  import torch
  from unit.common import DistributedTest
  from unit.simple_model import random_dataloader, SimpleModel
  from unit.util import bf16_required_version_check
  import deepspeed
  from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state
  from deepspeed.utils import safe_set_full_fp32_param, safe_set_full_optimizer_state
  from deepspeed.utils import safe_get_local_fp32_param, safe_get_local_grad, safe_get_local_optimizer_state
  from deepspeed.utils import safe_set_local_fp32_param, safe_set_local_optimizer_state
  from deepspeed.runtime.zero.offload_config import OffloadDeviceEnum
  from deepspeed.ops.aio import AsyncIOBuilder
label_desc: Use image node_id 1 to create a new MyModel object from inherited base classes: torch.nn with example: obj = MyModel(hidden_dim, frozen_weights)