class_docstr: ## Configurations
  
  This extends MNIST configurations to get the data loaders and Training and validation loop
  configurations to simplify our implementation.
class_docstr_tok: ["#","#","Configurations","This","extends","MNIST","configurations","to","get","the","data","loaders","and","Training","and","validation","loop","configurations","to","simplify","our","implementation","."]
docstr: Take a training step
docstr_tok: ["Take","a","training","step"]
code: def step(self, batch: Any, batch_idx: BatchIndex):
      """
      Take a training step
      """
  
      # Set model states
      self.generator.train(self.mode.is_train)
      self.discriminator.train(self.mode.is_train)
  
      # Get MNIST images
      data = batch[0].to(self.device)
  
      # Increment step in training mode
      if self.mode.is_train:
          tracker.add_global_step(len(data))
  
      # Train the discriminator
      with monit.section("discriminator"):
          # Get discriminator loss
          loss = self.calc_discriminator_loss(data)
  
          # Train
          if self.mode.is_train:
              self.discriminator_optimizer.zero_grad()
              loss.backward()
              if batch_idx.is_last:
                  tracker.add("discriminator", self.discriminator)
              self.discriminator_optimizer.step()
  
      # Train the generator once in every `discriminator_k`
      if batch_idx.is_interval(self.discriminator_k):
          with monit.section("generator"):
              loss = self.calc_generator_loss(data.shape[0])
  
              # Train
              if self.mode.is_train:
                  self.generator_optimizer.zero_grad()
                  loss.backward()
                  if batch_idx.is_last:
                      tracker.add("generator", self.generator)
                  self.generator_optimizer.step()
  
      tracker.save()
  
imports: from typing import Any
  import torch
  import torch.nn
  import torch.utils.data
  from torchvision import transforms
  from labml import tracker, monit, experiment
  from labml.configs import option, calculate
  from labml_helpers.datasets.mnist import MNISTConfigs
  from labml_helpers.device import DeviceConfigs
  from labml_helpers.module import Module
  from labml_helpers.optimizer import OptimizerConfigs
  from labml_helpers.train_valid import TrainValidConfigs, hook_model_outputs, BatchIndex
  from labml_nn.gan.original import DiscriminatorLogitsLoss, GeneratorLogitsLoss
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