class_docstr: ## Configurations
  
  This extends from [`CIFAR10Configs`](../experiments/cifar10.html) which defines all the
  dataset related configurations, optimizer, and a training loop.
class_docstr_tok: ["#","#","Configurations","This","extends","from","[","`","CIFAR10Configs","`","]","(","..","\/experiments\/cifar10.html",")","which","defines","all","the","dataset","related","configurations",",","optimizer",",","and","a","training","loop","."]
docstr: ### Training/validation step
  
  We define a custom training/validation step to include the distillation
docstr_tok: ["#","#","#","Training\/validation","step","We","define","a","custom","training\/validation","step","to","include","the","distillation"]
code: def step(self, batch: any, batch_idx: BatchIndex):
      """
      ### Training/validation step
  
      We define a custom training/validation step to include the distillation
      """
  
      # Training/Evaluation mode for the small model
      self.model.train(self.mode.is_train)
      # Large model in evaluation mode
      self.large.eval()
  
      # Move data to the device
      data, target = batch[0].to(self.device), batch[1].to(self.device)
  
      # Update global step (number of samples processed) when in training mode
      if self.mode.is_train:
          tracker.add_global_step(len(data))
  
      # Get the output logits, $v_i$, from the large model
      with torch.no_grad():
          large_logits = self.large(data)
  
      # Get the output logits, $z_i$, from the small model
      output = self.model(data)
  
      # Soft targets
      # $$p_i = \frac{\exp (\frac{v_i}{T})}{\sum_j \exp (\frac{v_j}{T})}$$
      soft_targets = nn.functional.log_softmax(
          large_logits / self.temperature, dim=-1
      )
      # Temperature adjusted probabilities of the small model
      # $$q_i = \frac{\exp (\frac{z_i}{T})}{\sum_j \exp (\frac{z_j}{T})}$$
      soft_prob = nn.functional.log_softmax(
          output / self.temperature, dim=-1
      )
  
      # Calculate the soft targets loss
      soft_targets_loss = self.kl_div_loss(soft_prob, soft_targets)
      # Calculate the true label loss
      label_loss = self.loss_func(output, target)
      # Weighted sum of the two losses
      loss = (
          self.soft_targets_weight * soft_targets_loss
          + self.label_loss_weight * label_loss
      )
      # Log the losses
      tracker.add(
          {
              "loss.kl_div.": soft_targets_loss,
              "loss.nll": label_loss,
              "loss.": loss,
          }
      )
  
      # Calculate and log accuracy
      self.accuracy(output, target)
      self.accuracy.track()
  
      # Train the model
      if self.mode.is_train:
          # Calculate gradients
          loss.backward()
          # 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/distillation/__init__.py
imports: import torch
  import torch.nn.functional
  from torch import nn
  from labml import experiment, tracker
  from labml.configs import option
  from labml_helpers.train_valid import BatchIndex
  from labml_nn.distillation.large import LargeModel
  from labml_nn.distillation.small import SmallModel
  from labml_nn.experiments.cifar10 import CIFAR10Configs
label_desc: Use image node_id 1 for calling the Configs obj's underlying member method code with example usage: obj.step(batch, batch_idx) without return types