code: def __init__(self, input_shape):
      super(MNIST_MLP, self).__init__()
      self.hdim = hdim = 390
      self.encoder = nn.Sequential(
          nn.Linear(
              input_shape[0] * input_shape[1] * input_shape[2], hdim
          ),
          nn.ReLU(True),
          nn.Linear(hdim, hdim),
          nn.ReLU(True),
      )
      self.n_outputs = hdim
  
      for m in self.encoder:
          if isinstance(m, nn.Linear):
              gain = nn.init.calculate_gain("relu")
              nn.init.xavier_uniform_(m.weight, gain=gain)
              nn.init.zeros_(m.bias)
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  import torchvision.models
label_desc: Use image node_id 1 to create a new MNIST_MLP object from inherited base classes: nn with example: obj = MNIST_MLP(input_shape)