class_docstr: ## Configurations
  
  This extends [`NLPAutoRegressionConfigs`](../../experiments/nlp_autoregression.html).
  
  The default configs can and will be over-ridden when we start the experiment
class_docstr_tok: ["#","#","Configurations","This","extends","[","`","NLPAutoRegressionConfigs","`","]","(","..","\/","..","\/experiments\/nlp_autoregression.html",")",".","The","default","configs","can","and","will","be","over-ridden","when","we","start","the","experiment"]
docstr: ### Training or validation step
docstr_tok: ["#","#","#","Training","or","validation","step"]
code: def step(self, batch: any, batch_idx: BatchIndex):
      """
      ### Training or validation step
      """
  
      # Move data to the device
      data, target = batch[0].to(self.device), batch[1].to(self.device)
  
      # Update global step (number of tokens processed) when in training mode
      if self.mode.is_train:
          tracker.add_global_step(data.shape[0] * data.shape[1])
  
      # Whether to capture model outputs
      with self.mode.update(is_log_activations=batch_idx.is_last):
          # Get model outputs.
          (
              output,
              counts,
              route_prob,
              n_dropped,
              route_prob_max,
          ) = self.model(data)
  
      # Calculate and cross entropy loss
      cross_entropy_loss = self.loss_func(output, target)
      # Total number of tokens processed, $T$, in the current batch $\mathscr{B}$
      total = counts.sum(dim=-1, keepdims=True)
      # Fraction of tokens routed to each expert
      # $$f_i = \frac{1}{T} \sum_{x \in \mathscr{B}} \mathbf{1} \{ \mathop{argmax} p(x), i \}$$
      # $f_i$ is the count of tokens where the argmax of $p(x)$ is equal to $i$.
      route_frac = counts / total
      # Mean routing probability
      # $$P_i = \frac{1}{T} \sum_{x \in \mathscr{B}} p_i (x)$$
      route_prob = route_prob / total
      # Load balancing loss
      # $$\mathscr{L} = N \sum_{i=1}^N f_i \cdot P_i$$
      # $\mathscr{L}$ is the loss for a single layer and here we are
      # taking the sum of losses across all layers.
      load_balancing_loss = (
          self.n_experts * (route_frac * route_prob).sum()
      )
  
      # Track stats
      tracker.add("dropped.", total.new_tensor(n_dropped) / total)
      tracker.add("route.min.", route_frac.min())
      tracker.add("route.max.", route_frac.max())
      tracker.add("route.std.", route_frac.std())
      tracker.add("route.max_prob.", route_prob_max)
      tracker.add("loss.", cross_entropy_loss)
      tracker.add("lb_loss.", load_balancing_loss)
  
      # Combined loss.
      # The load balancing loss is multiplied by a coefficient $\alpha$ which is
      # set to something small like $\alpha = 0.01$.
      loss = (
          cross_entropy_loss
          + self.load_balancing_loss_ceof * load_balancing_loss
      )
  
      # Calculate and log accuracy
      self.accuracy(output, target)
      self.accuracy.track()
  
      # Train the model
      if self.mode.is_train:
          # Calculate gradients
          loss.backward()
          # Clip gradients
          torch.nn.utils.clip_grad_norm_(
              self.model.parameters(), max_norm=self.grad_norm_clip
          )
          # Take optimizer step
          self.optimizer.step()
          # Log the model parameters and gradients on last batch of every epoch
          if batch_idx.is_last:
              tracker.add("model", self.model)
          # Clear the gradients
          self.optimizer.zero_grad()
  
      # Save the tracked metrics
      tracker.save()
  
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 Configs obj's underlying member method code with example usage: obj.step(batch, batch_idx) without return types