class_docstr: This is a configurable module that works for many standard DL experiments.
  This is based on :class:`labml_helpers.training_loop.TrainValidConfigs`.
  
  Arguments:
      model: A PyTorch model.
      optimizer: A PyTorch optimizer to update model.
      device: The device to train the model on. This defaults to a configurable device -
       :class:`labml_helpers.device.DeviceConfigs`.
      loss_function: A function to calculate the loss. This should accept ``model_output, target`` as
       arguments.
      update_batches (int): Number of batches to accumulate before taking an optimizer step.
       Defaults to ``1``.
      log_params_updates (int): How often (number of batches) to track model parameters and gradients.
       Defaults to a large number; i.e. logs every epoch.
      log_activations_batches (int): How often to log model activations. 
       Defaults to a large number; i.e. logs every epoch.
      log_save_batches (int): How often to call :func:`labml.tracker.save`.
class_docstr_tok: ["This","is","a","configurable","module","that","works","for","many","standard","DL","experiments",".","This","is","based","on",":","class",":","`","labml_helpers.training_loop.TrainValidConfigs","`",".","Arguments",":","model",":","A","PyTorch","model",".","optimizer",":","A","PyTorch","optimizer","to","update","model",".","device",":","The","device","to","train","the","model","on",".","This","defaults","to","a","configurable","device","-",":","class",":","`","labml_helpers.device.DeviceConfigs","`",".","loss_function",":","A","function","to","calculate","the","loss",".","This","should","accept","``","model_output",",","target","``","as","arguments",".","update_batches","(","int",")",":","Number","of","batches","to","accumulate","before","taking","an","optimizer","step",".","Defaults","to","``","1","``",".","log_params_updates","(","int",")",":","How","often","(","number","of","batches",")","to","track","model","parameters","and","gradients",".","Defaults","to","a","large","number",";","i.e",".","logs","every","epoch",".","log_activations_batches","(","int",")",":","How","often","to","log","model","activations",".","Defaults","to","a","large","number",";","i.e",".","logs","every","epoch",".","log_save_batches","(","int",")",":","How","often","to","call",":","func",":","`","labml.tracker.save","`","."]
code: def step(self, batch: Any, batch_idx: BatchIndex):
      self.model.train(self.mode.is_train)
      data, target = batch[0].to(self.device), batch[1].to(self.device)
  
      if self.mode.is_train:
          tracker.add_global_step(len(data))
  
      is_log_activations = batch_idx.is_interval(
          self.log_activations_batches
      )
      with monit.section("model"):
          with self.mode.update(is_log_activations=is_log_activations):
              output = self.model(data)
  
      loss = self.loss_func(output, target)
      tracker.add("loss.", loss)
  
      if self.mode.is_train:
          with monit.section("backward"):
              loss.backward()
  
          if batch_idx.is_interval(self.update_batches):
              with monit.section("optimize"):
                  self.optimizer.step()
              if batch_idx.is_interval(self.log_params_updates):
                  tracker.add("model", self.model)
              self.optimizer.zero_grad()
  
          if batch_idx.is_interval(self.log_save_batches):
              tracker.save()
  
imports: from typing import Optional, Dict, List, Callable, Any
  import torch.optim
  import torch.optim
  import torch.utils.data
  import torch.utils.data
  from torch import nn
  import labml.utils.pytorch
  from labml import tracker, monit
  from labml.configs import option, meta_config
  from .device import DeviceConfigs
  from .metrics import StateModule
  from .training_loop import TrainingLoopConfigs
label_desc: Use image node_id 2 for calling the SimpleTrainValidConfigs obj's underlying member method code with example usage: obj.step(batch, batch_idx) without return types