code: def __init__(
      self,
      num_experts: int,
      num_experts_per_tok: int,
      num_shards: int,
      gate_softmax: bool = False,
      **kwargs,
  ):
      super().__init__()
      self.experts = nn.ModuleList(
          [
              TorchFFN(**kwargs).to(f"cuda:{i//num_shards}")
              for i in range(num_experts)
          ]
      )
      self.gate = nn.Linear(kwargs["dim"], num_experts, bias=False)
  
      self.num_experts_per_tok = num_experts_per_tok
      self.gate_softmax = gate_softmax
      print("Softmax for Gate:{}".format(str(gate_softmax)))
  
imports: import math
  from dataclasses import dataclass
  from typing import Dict, Optional, Tuple
  import torch
  import torch.nn.functional
  from torch import nn
  from .utils import ModelArgs
  from .attention import TorchAttention, FairScaleAttention
  from .ffn import TorchFFN, FairScaleFFN
  from .transformer import TorchTransformerBlock, TorchTransformer, FairScaleTransformer
label_desc: Use image node_id 1 to create a new MoETorchFFN object from inherited base classes: nn with example: obj = MoETorchFFN(num_experts, num_experts_per_tok, num_shards, gate_softmax)