class_docstr: ## Auto regressive model
class_docstr_tok: ["#","#","Auto","regressive","model"]
code: def forward(self, x: torch.Tensor):
      # Initialize the subsequent mask
      if self.mask is None or self.mask.size(0) != len(x):
          from labml_nn.transformers.utils import subsequent_mask
  
          self.mask = subsequent_mask(len(x)).to(x.device)
      # Token embeddings
      x = self.src_embed(x)
      # Run it through the transformer
      (
          res,
          counts,
          route_prob,
          n_dropped,
          route_prob_max,
      ) = self.transformer(x, self.mask)
      # Generate logits of the next token
      res = self.generator(res)
      #
      return res, counts, route_prob, n_dropped, route_prob_max
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/switch/experiment.py
imports: import torch
  import torch.nn
  from labml import experiment, tracker
  from labml.configs import option
  from labml_helpers.module import Module
  from labml_helpers.train_valid import BatchIndex
  from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
label_desc: Use image node_id 2 for calling the AutoregressiveModel obj's underlying member method code with example usage: obj.forward(x) and returns: res,  counts,  route_prob,  n_dropped,  route_prob_max