code: def mnist_net(num_classes=10):
      return nn.Sequential(
          OrderedDict(
              [
                  ("conv1", nn.Conv2d(1, 32, 5, padding="same")),
                  ("relu1", nn.ReLU()),
                  ("pool1", nn.MaxPool2d(2, stride=2, padding=0)),
                  ("bn1", nn.BatchNorm2d(32, eps=0.001, momentum=0.01)),
                  ("conv2", nn.Conv2d(32, 64, 5, padding="same")),
                  ("relu2", nn.ReLU()),
                  ("pool2", nn.MaxPool2d(2, stride=2, padding=0)),
                  ("flatten", nn.Flatten()),
                  ("dense1", nn.Linear(3136, 1024)),
                  ("relu3", nn.ReLU()),
                  ("dropout", nn.Dropout(p=0.4)),
                  ("dense2", nn.Linear(1024, num_classes)),
                  ("softmax", nn.LogSoftmax()),
              ]
          )
      )
  
file_path: coremltools/docs/examples/optimize/torch/palettization/dkm_palettization.py
imports: from collections import OrderedDict
  import torch
  import torch.nn
  import torch.nn.functional
  import os
  from torchvision import datasets, transforms
  from coremltools.optimize.torch.palettization import DKMPalettizer, DKMPalettizerConfig
  import coremltools
label_desc: Use image node_id 1 for calling a global function with example usage: mnist_net(num_classes) and returns: nn