class_docstr: Call a Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts.
  Args:
  input_size: integer - size of the input
  output_size: integer - size of the input
  num_experts: an integer - number of experts
  hidden_size: an integer - hidden size of the experts
  noisy_gating: a boolean
  k: an integer - how many experts to use for each batch element
docstr: Args:
  x: tensor shape [batch_size, input_size]
  train: a boolean scalar.
  loss_coef: a scalar - multiplier on load-balancing losses
  
  Returns:
  y: a tensor with shape [batch_size, output_size].
  extra_training_loss: a scalar.  This should be added into the overall
  training loss of the model.  The backpropagation of this loss
  encourages all experts to be approximately equally used across a batch.
docstr_tok: ["Args",":","x",":","tensor","shape","[","batch_size",",","input_size","]","train",":","a","boolean","scalar",".","loss_coef",":","a","scalar","-","multiplier","on","load-balancing","losses","Returns",":","y",":","a","tensor","with","shape","[","batch_size",",","output_size","]",".","extra_training_loss",":","a","scalar",".","This","should","be","added","into","the","overall","training","loss","of","the","model",".","The","backpropagation","of","this","loss","encourages","all","experts","to","be","approximately","equally","used","across","a","batch","."]
code: def forward(self, x, loss_coef=1e-2):
      """Args:
      x: tensor shape [batch_size, input_size]
      train: a boolean scalar.
      loss_coef: a scalar - multiplier on load-balancing losses
  
      Returns:
      y: a tensor with shape [batch_size, output_size].
      extra_training_loss: a scalar.  This should be added into the overall
      training loss of the model.  The backpropagation of this loss
      encourages all experts to be approximately equally used across a batch.
      """
      gates, load = self.noisy_top_k_gating(x, self.training)
      # calculate importance loss
      importance = gates.sum(0)
      #
      loss = self.cv_squared(importance) + self.cv_squared(load)
      loss *= loss_coef
  
      dispatcher = SparseDispatcher(self.num_experts, gates)
      expert_inputs = dispatcher.dispatch(x)
      gates = dispatcher.expert_to_gates()
      expert_outputs = [
          self.experts[i](expert_inputs[i])
          for i in range(self.num_experts)
      ]
      y = dispatcher.combine(expert_outputs)
      return y, loss
  
imports: import torch
  import torch.nn
  from torch.distributions.normal import Normal
  import numpy
label_desc: Use image node_id 6 for calling the MoE obj's underlying member method code with example usage: obj.forward(x, loss_coef) and returns: y,  loss