class_docstr: MOELayer module which implements MixtureOfExperts as described in Gshard_.
  ::
  
      gate = TopKGate(model_dim, num_experts)
      moe = MOELayer(gate, expert)
      output = moe(input)
      l_aux = moe.l_aux
  
  .. Gshard_: https://arxiv.org/pdf/2006.16668.pdf
  
  Args:
      gate (torch.nn.Module):
          gate network
      expert (torch.nn.Module):
          expert network
class_docstr_tok: ["MOELayer","module","which","implements","MixtureOfExperts","as","described","in","Gshard_",".",":",":","gate","=","TopKGate","(","model_dim",",","num_experts",")","moe","=","MOELayer","(","gate",",","expert",")","output","=","moe","(","input",")","l_aux","=","moe.l_aux","..","Gshard_",":","https",":","\/\/arxiv.org\/pdf\/2006.16668.pdf","Args",":","gate","(","torch.nn.Module",")",":","gate","network","expert","(","torch.nn.Module",")",":","expert","network"]
code: def _set_ep_group(self, ep_group):
      self.ep_group = ep_group
  
imports: from deepspeed.utils.timer import SynchronizedWallClockTimer
  from deepspeed.utils import logger
  from typing import Callable, Dict, TYPE_CHECKING, Any, Optional, Tuple
  import torch
  from torch import Tensor
  from torch.nn import Module
  import torch.nn.functional
  from deepspeed.utils import groups
  from .mappings import drop_tokens, gather_tokens
  from deepspeed import comm
label_desc: Use image node_id 2 for calling the MOELayer obj's underlying member method code with example usage: obj._set_ep_group(ep_group) without return types