class_docstr: ## Switch Transformer
class_docstr_tok: ["#","#","Switch","Transformer"]
code: def forward(self, x: torch.Tensor, mask: torch.Tensor):
      # Run through each transformer layer
      counts, route_prob, n_dropped, route_prob_max = [], [], [], []
      for layer in self.layers:
          x, f, p, n_d, p_max = layer(x=x, mask=mask)
          counts.append(f)
          route_prob.append(p)
          n_dropped.append(n_d)
          route_prob_max.append(p_max)
      # Finally, normalize the vectors
      x = self.norm(x)
      #
      return (
          x,
          torch.stack(counts),
          torch.stack(route_prob),
          n_dropped,
          torch.stack(route_prob_max),
      )
  
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 2 for calling the SwitchTransformer obj's underlying member method code with example usage: obj.forward(x, mask) and returns: x,  torch,  torch,  n_dropped,  torch