code: def __init__(self, use_moe):
      super(Net, self).__init__()
      self.use_moe = use_moe
  
      if self.use_moe:
          self.moe_ffn = moe.moe_layer(
              gate_type={
                  "type": "top",
                  "k": 1,
                  "capacity_factor": 0,
                  "gate_noise": 1.0,
              },
              experts={
                  "type": "ffn",
                  "count_per_node": 1,
                  "hidden_size_per_expert": 128,
                  "output_dim": 10,
                  "activation_fn": lambda x: self.dropout2(F.relu(x)),
              },
              model_dim=9216,
              seeds=(1, penv.global_rank + 1),
              scan_expert_func=lambda name, param: setattr(
                  param, "skip_allreduce", True
              ),
          )
      else:
          torch.manual_seed(1)
          self.fc1 = nn.Linear(9216, 128)
          self.fc2 = nn.Linear(128, 10)
  
      torch.manual_seed(1)
      self.conv1 = nn.Conv2d(1, 32, 3, 1)
      self.conv2 = nn.Conv2d(32, 64, 3, 1)
      self.dropout1 = nn.Dropout(0.25)
      self.dropout2 = nn.Dropout(0.5)
  
file_path: tutel/tutel/examples/moe_mnist.py
imports: import argparse
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torchvision import datasets, transforms
  from torch.optim.lr_scheduler import StepLR
  from tutel import system
  from tutel import moe
  from tutel import net
  import logging
label_desc: Use image node_id 1 to create a new Net object from inherited base classes: nn with example: obj = Net(use_moe)