code: def forward(self, x):
      x = self.pool1(self.relu1(self.conv1(x)))
      x = self.pool2(self.relu2(self.conv2(x)))
      x = torch.flatten(x, 1)
      x = self.relu3(self.fc1(x))
      x = self.relu4(self.fc2(x))
      x = self.fc3(x)
      return F.log_softmax(x, dim=1)
  
file_path: nni/nni_assets/compression/mnist_model.py
imports: from pathlib import Path
  import torch
  from torch import nn
  from torch.nn import functional
  from torchvision import datasets, transforms
label_desc: Use image node_id 2 for calling the TorchModel obj's underlying member method code with example usage: obj.forward(x) and returns: F