class_docstr: ## Configurations
  
  The default configurations can and will be overridden when we start the experiment.
class_docstr_tok: ["#","#","Configurations","The","default","configurations","can","and","will","be","overridden","when","we","start","the","experiment","."]
docstr: ### Training/validation step
docstr_tok: ["#","#","#","Training\/validation","step"]
code: def step(self, batch: any, batch_idx: BatchIndex):
      """
      ### Training/validation step
      """
  
      # Move data to the device
      data, target = batch[0].to(self.device), batch[1].to(self.device)
  
      # Update global step (number of tokens processed) when in training mode
      if self.mode.is_train:
          tracker.add_global_step(data.shape[0] * data.shape[1])
  
      # Whether to capture model outputs
      with self.mode.update(is_log_activations=batch_idx.is_last):
          # Get memories
          mem = self.memory.get()
          # Run the model
          output, new_mem = self.model(data, mem)
          # Merge and compress memory
          mem, mem_to_compress = self.merge_compress_memory(
              mem, new_mem
          )
          # Update memories
          self.memory.set(mem)
  
      # Calculate and log cross entropy loss
      loss = self.loss_func(output, target)
      tracker.add("loss.", loss)
  
      # Calculate attention reconstruction loss if memories were compressed in this step
      if mem_to_compress:
          # Get attention reconstruction loss
          ar_loss = self.attention_reconstruction_loss(
              new_mem, mem_to_compress
          )
          # Track attention reconstruction loss
          tracker.add("ar_loss.", ar_loss)
          # Add attention reconstruction loss to loss
          loss = loss + ar_loss
  
      # Calculate and log accuracy
      self.accuracy(output, target)
      self.accuracy.track()
  
      # Train the model
      if self.mode.is_train:
          # Calculate gradients
          loss.backward()
          # Clip gradients
          torch.nn.utils.clip_grad_norm_(
              self.model.parameters(), max_norm=self.grad_norm_clip
          )
          # Take optimizer step
          self.optimizer.step()
          # Log the model parameters and gradients on last batch of every epoch
          if batch_idx.is_last:
              tracker.add("model", self.model)
          # Clear the gradients
          self.optimizer.zero_grad()
  
      # Save the tracked metrics
      tracker.save()
  
imports: from typing import List, Tuple, NamedTuple
  import torch
  import torch.nn
  from labml import experiment, tracker, monit, logger
  from labml.configs import option
  from labml.logger import Text
  from labml_helpers.metrics.simple_state import SimpleStateModule
  from labml_helpers.module import Module
  from labml_helpers.train_valid import BatchIndex, hook_model_outputs
  from labml_nn.experiments.nlp_autoregression import NLPAutoRegressionConfigs
  from labml_nn.transformers.compressive import CompressiveTransformer, AttentionReconstructionLoss, CompressiveTransformerLayer, Conv1dCompression
label_desc: Use image node_id 3 for calling the Configs obj's underlying member method code with example usage: obj.step(batch, batch_idx) without return types