class_docstr: ## Configurable Experiment Definition
code: def step(self, batch: any, batch_idx: BatchIndex):
      # Get the batch
      data, target = batch[0].to(self.device), batch[1].to(self.device)
  
      # Add global step if we are in training mode
      if self.mode.is_train:
          tracker.add_global_step(len(data))
  
      # Run the model and specify whether to log the activations
      with self.mode.update(is_log_activations=batch_idx.is_last):
          output = self.model(data)
  
      # Calculate the loss
      loss = self.loss_func(output, target)
      # Calculate the accuracy
      self.accuracy_func(output, target)
      # Log the loss
      tracker.add("loss.", loss)
  
      # Optimize if we are in training mode
      if self.mode.is_train:
          # Calculate the gradients
          loss.backward()
  
          # Take optimizer step
          self.optimizer.step()
          # Log the parameter and gradient L2 norms once per epoch
          if batch_idx.is_last:
              tracker.add("model", self.model)
              tracker.add(
                  "optimizer", (self.optimizer, {"model": self.model})
              )
          # Clear the gradients
          self.optimizer.zero_grad()
  
      # Save logs
      tracker.save()
  
imports: import torch.nn
  import torch.utils.data
  from labml_helpers.module import Module
  from labml import experiment, tracker
  from labml.configs import option
  from labml_helpers.datasets.mnist import MNISTConfigs
  from labml_helpers.device import DeviceConfigs
  from labml_helpers.metrics.accuracy import Accuracy
  from labml_helpers.seed import SeedConfigs
  from labml_helpers.train_valid import TrainValidConfigs, BatchIndex, hook_model_outputs
  from labml_nn.optimizers.configs import OptimizerConfigs
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