class_docstr: Top-1 MoE
  
  Curently supports Sinkhorn based expert routing.
class_docstr_tok: ["Top-1","MoE","Curently","supports","Sinkhorn","based","expert","routing","."]
code: def forward(self, hidden_states):
      hidden_shape = hidden_states.shape
      route, _ = self.router(hidden_states)
      route = route.view(-1, self.num_experts)
      if self.training:
          with torch.no_grad():
              norm_route = self.route_algo(
                  route.detach().to(dtype=torch.float32)
              )  # explicit fp32 conversion for stability
              _, max_ind = torch.max(norm_route, dim=1)
          route = torch.sigmoid(route)
          max_prob = route[torch.arange(route.size(0)), max_ind]
      else:
          route = torch.sigmoid(route)
          max_prob, max_ind = torch.max(route, dim=1)
      max_prob = torch.unsqueeze(max_prob, 1)
  
      hidden_states = hidden_states.view(-1, hidden_shape[-1])
  
      local_indices = (max_ind == 0).nonzero()
      hidden = hidden_states[local_indices, :]
      output, output_bias = self.experts[0](hidden)
      output_bias = output_bias.expand_as(output)
  
      output_total = torch.empty_like(hidden_states, dtype=output.dtype)
      output_bias_total = torch.empty_like(
          hidden_states, dtype=output_bias.dtype
      )
  
      output_total[local_indices, :] = output
      output_bias_total[local_indices, :] = output_bias
  
      for expert_num, expert in enumerate(self.experts):
          if expert_num == 0:
              continue
          local_indices = (max_ind == expert_num).nonzero()
          hidden = hidden_states[local_indices, :]
          output, output_bias = expert(hidden)
          output_bias = output_bias.expand_as(output)
          output_total[local_indices, :] = output
          output_bias_total[local_indices, :] = output_bias
  
      output_total = output_total * max_prob
      output_bias_total = output_bias_total * max_prob
      output_total = output_total.view(hidden_shape)
      output_bias_total = output_bias_total.view(hidden_shape)
  
      return output_total, output_bias_total
  
imports: import torch
  import torch.nn.functional
  from nemo.collections.nlp.modules.common.megatron.adapters.parallel_adapters import AdapterName, MLPInfusedAdapterConfig
  from nemo.collections.nlp.modules.common.megatron.fused_bias_geglu import fused_bias_geglu
  from nemo.collections.nlp.modules.common.megatron.fused_bias_gelu import fused_bias_gelu
  from nemo.collections.nlp.modules.common.megatron.fused_layer_norm import get_layer_norm
  from nemo.collections.nlp.modules.common.megatron.layer_norm_1p import LayerNorm1P
  from nemo.collections.nlp.modules.common.megatron.module import MegatronModule
  from nemo.collections.nlp.modules.common.megatron.utils import ApexGuardDefaults, ApproxGELUActivation, erf_gelu
  from nemo.collections.nlp.modules.common.megatron.utils import openai_gelu
  from nemo.collections.nlp.modules.common.megatron.utils import squared_relu
  from nemo.core import adapter_mixins
label_desc: Use image node_id 2 for calling the SwitchMLP obj's underlying member method code with example usage: obj.forward(hidden_states) and returns: output_total,  output_bias_total