code: def test_mnist_model():
      class DepthwiseSeparableConv(nn.Module):
          def __init__(self, in_ch, out_ch):
              super().__init__()
              self.depthwise = nn.Conv2d(
                  in_ch, in_ch, kernel_size=3, groups=in_ch
              )
              self.pointwise = nn.Conv2d(in_ch, out_ch, kernel_size=1)
  
          def forward(self, x):
              return self.pointwise(self.depthwise(x))
  
      class Net(ModelSpace):
          def __init__(self):
              super().__init__()
              self.conv1 = nn.Conv2d(1, 32, 3, 1)
              self.conv2 = LayerChoice(
                  [
                      nn.Conv2d(32, 64, 3, 1),
                      DepthwiseSeparableConv(32, 64),
                  ],
                  label="conv",
              )
              self.pool = nn.MaxPool2d(2)
              self.dropout1 = MutableDropout(
                  nni.choice("dropout", [0.25, 0.5, 0.75])
              )
              self.dropout2 = nn.Dropout(0.5)
              feature = nni.choice("feature", [64, 128, 256])
              self.fc1 = MutableLinear(9216, feature)
              self.fc2 = MutableLinear(feature, 10)
  
          def forward(self, x):
              x = F.relu(self.conv1(x))
              x = self.pool(self.conv2(x))
              x = self.dropout1(x).view(x.size(0), -1)
              x = self.fc2(self.dropout2(F.relu(self.fc1(x))))
              return x
  
      input = ShapeTensor(torch.randn(4, 1, 28, 28), True)
      output = shape_inference(Net(), input)
      assert output.real_shape == MutableShape(4, 10)
  
      net = Net()
      shapes = submodule_input_output_shapes(net, input)
      input_shapes = {
          k: v[0][0] for k, v in shapes.items()
      }  # unwrap tuple
      output_shapes = {k: v[1] for k, v in shapes.items()}
      expected_input_shapes = {
          "": MutableShape(4, 1, 28, 28),
          "conv1": MutableShape(4, 1, 28, 28),
          "conv2": MutableShape(4, 32, 26, 26),
          "conv2.0": MutableShape(4, 32, 26, 26),
          "conv2.1": MutableShape(4, 32, 26, 26),
          "conv2.1.depthwise": MutableShape(4, 32, 26, 26),
          "conv2.1.pointwise": MutableShape(4, 32, 24, 24),
          "dropout1": MutableShape(4, 64, 12, 12),
          "dropout2": MutableShape(
              4,
              Categorical(
                  [64, 128, 256],
                  label=net.fc1.args["out_features"].label,
              ),
          ),
          "fc1": MutableShape(4, 9216),
          "fc2": MutableShape(
              4,
              Categorical(
                  [64, 128, 256],
                  label=net.fc1.args["out_features"].label,
              ),
          ),
          "pool": MutableShape(4, 64, 24, 24),
      }
      expected_output_shapes = {
          "": MutableShape(4, 10),
          "conv1": MutableShape(4, 32, 26, 26),
          "conv2": MutableShape(4, 64, 24, 24),
          "conv2.0": MutableShape(4, 64, 24, 24),
          "conv2.1": MutableShape(4, 64, 24, 24),
          "conv2.1.depthwise": MutableShape(4, 32, 24, 24),
          "conv2.1.pointwise": MutableShape(4, 64, 24, 24),
          "dropout1": MutableShape(4, 64, 12, 12),
          "dropout2": MutableShape(
              4,
              Categorical(
                  [64, 128, 256],
                  label=net.fc1.args["out_features"].label,
              ),
          ),
          "fc1": MutableShape(
              4,
              Categorical(
                  [64, 128, 256],
                  label=net.fc1.args["out_features"].label,
              ),
          ),
          "fc2": MutableShape(4, 10),
          "pool": MutableShape(4, 64, 12, 12),
      }
      assert input_shapes == expected_input_shapes
      assert output_shapes == expected_output_shapes
  
imports: import torch
  import torch.nn.functional
  from torch import nn
  import pytest
  import nni
  import nni.nas.profiler.pytorch.utils.shape
  from nni.mutable import MutableList, Categorical
  from nni.nas.nn.pytorch import MutableLinear, MutableConv2d, MutableDropout, MutableMultiheadAttention, ModelSpace, LayerChoice, InputChoice, Repeat
  from nni.nas.profiler.pytorch.utils import shape_inference, submodule_input_output_shapes, ShapeTensor, MutableShape, profiler_leaf_module
label_desc: Use image node_id 27 for calling a global function with example usage: test_mnist_model() and returns: self, x