class_docstr: # Switch Transformer Block
  
  This is the same as [normal transformer block](../models.html#TransformerLayer)
  with handling extra outputs of switch feedforward module.
class_docstr_tok: ["#","Switch","Transformer","Block","This","is","the","same","as","[","normal","transformer","block","]","(","..","\/models.html","#","TransformerLayer",")","with","handling","extra","outputs","of","switch","feedforward","module","."]
docstr: * `d_model` is the token embedding size
  * `attn` is the attention module
  * `feed_forward` is the feed forward module (which is the switching module in this case)
  * `dropout_prob` is the probability of dropping out after self attention and FFN
docstr_tok: ["*","`","d_model","`","is","the","token","embedding","size","*","`","attn","`","is","the","attention","module","*","`","feed_forward","`","is","the","feed","forward","module","(","which","is","the","switching","module","in","this","case",")","*","`","dropout_prob","`","is","the","probability","of","dropping","out","after","self","attention","and","FFN"]
code: def __init__(
      self,
      *,
      d_model: int,
      attn: MultiHeadAttention,
      feed_forward: SwitchFeedForward,
      dropout_prob: float,
  ):
      """
      * `d_model` is the token embedding size
      * `attn` is the attention module
      * `feed_forward` is the feed forward module (which is the switching module in this case)
      * `dropout_prob` is the probability of dropping out after self attention and FFN
      """
      super().__init__()
      self.size = d_model
      self.attn = attn
      self.feed_forward = feed_forward
      self.dropout = nn.Dropout(dropout_prob)
      self.norm_self_attn = nn.LayerNorm([d_model])
      self.norm_ff = nn.LayerNorm([d_model])
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/switch/__init__.py
imports: import torch
  from torch import nn
  from labml_helpers.module import Module
  from labml_nn.transformers.feed_forward import FeedForward
  from labml_nn.transformers.mha import MultiHeadAttention
  from labml_nn.utils import clone_module_list
label_desc: Use image node_id 1 to create a new SwitchTransformerLayer object from inherited base classes: Module with example: obj = SwitchTransformerLayer()