class_docstr: ## Configurations
  
  We use [`MNISTConfigs`](../../experiments/mnist.html#MNISTConfigs) configurations.
class_docstr_tok: ["#","#","Configurations","We","use","[","`","MNISTConfigs","`","]","(","..","\/","..","\/experiments\/mnist.html","#","MNISTConfigs",")","configurations","."]
docstr: ### Training or validation step
docstr_tok: ["#","#","#","Training","or","validation","step"]
code: def step(self, batch: Any, batch_idx: BatchIndex):
      """
      ### Training or validation step
      """
  
      # Training/Evaluation mode
      self.model.train(self.mode.is_train)
  
      # Move data to the device
      data, target = batch[0].to(self.device), batch[1].to(self.device)
  
      # One-hot coded targets
      eye = torch.eye(10).to(torch.float).to(self.device)
      target = eye[target]
  
      # Update global step (number of samples processed) when in training mode
      if self.mode.is_train:
          tracker.add_global_step(len(data))
  
      # Get model outputs
      outputs = self.model(data)
      # Get evidences $e_k \ge 0$
      evidence = self.outputs_to_evidence(outputs)
  
      # Calculate loss
      loss = self.loss_func(evidence, target)
      # Calculate KL Divergence regularization loss
      kl_div_loss = self.kl_div_loss(evidence, target)
      tracker.add("loss.", loss)
      tracker.add("kl_div_loss.", kl_div_loss)
  
      # KL Divergence loss coefficient $\lambda_t$
      annealing_coef = min(
          1.0, self.kl_div_coef(tracker.get_global_step())
      )
      tracker.add("annealing_coef.", annealing_coef)
  
      # Total loss
      loss = loss + annealing_coef * kl_div_loss
  
      # Track statistics
      self.stats(evidence, target)
  
      # Train the model
      if self.mode.is_train:
          # Calculate gradients
          loss.backward()
          # Take optimizer step
          self.optimizer.step()
          # Clear the gradients
          self.optimizer.zero_grad()
  
      # Save the tracked metrics
      tracker.save()
  
imports: from typing import Any
  import torch.nn
  import torch.utils.data
  from labml import tracker, experiment
  from labml.configs import option, calculate
  from labml_helpers.module import Module
  from labml_helpers.schedule import Schedule, RelativePiecewise
  from labml_helpers.train_valid import BatchIndex
  from labml_nn.experiments.mnist import MNISTConfigs
  from labml_nn.uncertainty.evidence import KLDivergenceLoss, TrackStatistics, MaximumLikelihoodLoss, CrossEntropyBayesRisk, SquaredErrorBayesRisk
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