class_docstr: Configurations with a
   [simple training loop](https://docs.labml.ai/api/helpers.html#labml_helpers.train_valid.SimpleTrainValidConfigs)
docstr: This method gets called by the trainer for each batch
docstr_tok: ["This","method","gets","called","by","the","trainer","for","each","batch"]
code: def step(self, batch: Any, batch_idx: BatchIndex):
      """
      This method gets called by the trainer for each batch
      """
      # Set the model mode
      self.model.train(self.mode.is_train)
  
      # Get the input 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))
  
      # Run the model
      p, y_hat, p_sampled, y_hat_sampled = self.model(data)
  
      # Calculate the reconstruction loss
      loss_rec = self.loss_rec(p, y_hat, target.to(torch.float))
      tracker.add("loss.", loss_rec)
  
      # Calculate the regularization loss
      loss_reg = self.loss_reg(p)
      tracker.add("loss_reg.", loss_reg)
  
      # $L = L_{Rec} + \beta L_{Reg}$
      loss = loss_rec + self.beta * loss_reg
  
      # Calculate the expected number of steps taken
      steps = torch.arange(1, p.shape[0] + 1, device=p.device)
      expected_steps = (p * steps[:, None]).sum(dim=0)
      tracker.add("steps.", expected_steps)
  
      # Call accuracy metric
      self.accuracy(y_hat_sampled > 0, target)
  
      if self.mode.is_train:
          # Compute gradients
          loss.backward()
          # Clip gradients
          torch.nn.utils.clip_grad_norm_(
              self.model.parameters(), max_norm=self.grad_norm_clip
          )
          # Optimizer
          self.optimizer.step()
          # Clear gradients
          self.optimizer.zero_grad()
          #
          tracker.save()
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/adaptive_computation/ponder_net/experiment.py
imports: from typing import Any
  import torch
  from torch import nn
  from torch.utils.data import DataLoader
  from labml import tracker, experiment
  from labml_helpers.metrics.accuracy import AccuracyDirect
  from labml_helpers.train_valid import SimpleTrainValidConfigs, BatchIndex
  from labml_nn.adaptive_computation.parity import ParityDataset
  from labml_nn.adaptive_computation.ponder_net import ParityPonderGRU, ReconstructionLoss, RegularizationLoss
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