class_docstr: Configurations with MNIST data and Train & Validation setup
class_docstr_tok: ["Configurations","with","MNIST","data","and","Train","&","Validation","setup"]
docstr: This method gets called by the trainer
docstr_tok: ["This","method","gets","called","by","the","trainer"]
code: def step(self, batch: Any, batch_idx: BatchIndex):
      """
      This method gets called by the trainer
      """
      # Set the model mode
      self.model.train(self.mode.is_train)
  
      # Get the images and labels and move them to the model's device
      data, target = batch[0].to(self.device), batch[1].to(self.device)
  
      # Increment step in training mode
      if self.mode.is_train:
          tracker.add_global_step(len(data))
  
      # Whether to log activations
      with self.mode.update(is_log_activations=batch_idx.is_last):
          # Run the model
          caps, reconstructions, pred = self.model(data)
  
      # Calculate the total loss
      loss = self.margin_loss(
          caps, target
      ) + 0.0005 * self.reconstruction_loss(reconstructions, data)
      tracker.add("loss.", loss)
  
      # Call accuracy metric
      self.accuracy(pred, target)
  
      if self.mode.is_train:
          loss.backward()
  
          self.optimizer.step()
          # Log parameters and gradients
          if batch_idx.is_last:
              tracker.add("model", self.model)
          self.optimizer.zero_grad()
  
          tracker.save()
  
imports: from typing import Any
  import torch.nn
  import torch.nn.functional
  import torch.utils.data
  from labml import experiment, tracker
  from labml.configs import option
  from labml_helpers.datasets.mnist import MNISTConfigs
  from labml_helpers.metrics.accuracy import AccuracyDirect
  from labml_helpers.module import Module
  from labml_helpers.train_valid import SimpleTrainValidConfigs, BatchIndex
  from labml_nn.capsule_networks import Squash, Router, MarginLoss
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