code: def training_step(self, batch, batch_idx):
      output: HyperpriorOutput
      opts = self.optimizers()
      scheds = self.lr_schedulers()
      if len(opts) == 3:
          model_opt, disc_opt, quantile_opt = opts
          model_sched, disc_sched, quantile_sched = scheds
      else:
          model_opt, disc_opt = opts
          model_sched, disc_sched = scheds
          quantile_opt = None
  
      # if doing Monte Carlo sampling, split the batch for the discriminator
      model_images, disc_images = self._resample_batch(
          batch, mc_sampling=self.mc_sampling
      )
  
      # get the discriminator/generator targets
      if self.latent_projector is not None:
          with torch.no_grad():
              fake_target = self.latent_projector(model_images)
              real_target = self.latent_projector(disc_images)
      else:
          real_target = torch.tensor(1, dtype=torch.long)
          fake_target = torch.tensor(1, dtype=torch.long)
  
      ########### GENERATOR STEP ###########
      self.disable_gradients(disc_opt)
      self.enable_gradients(model_opt)
  
      output = self.model(model_images)
  
      # distortion loss
      distortion_loss = self.distortion_loss(output.image, model_images)
  
      # generator loss
      g_loss = self.generator_loss(
          self.run_discriminator(output.image, output.latent.detach()),
          fake_target,
      )
  
      # rate loss
      rate_output = self.rate_loss(
          original=model_images,
          latent_likelihoods=output.latent_likelihoods,
          quantized_latent_likelihoods=output.quantized_latent_likelihoods,
          hyper_latent_likelihoods=output.hyper_latent_likelihoods,
          quantized_hyper_latent_likelihoods=output.quantized_hyper_latent_likelihoods,
      )
  
      # total loss
      model_loss = (
          self.distortion_lam * distortion_loss
          + rate_output.rate_loss
          + self.generator_weight * g_loss
      )
  
      # step optimizer
      model_opt.zero_grad()
      self.manual_backward(model_loss)
      if self.gradient_clip_val is not None:
          self.clip_gradients(
              model_opt,
              gradient_clip_val=self.gradient_clip_val,
              gradient_clip_algorithm=self.gradient_clip_algorithm,
          )
      model_opt.step()
      model_sched.step()
  
      ########### DISCRIMINATOR STEP ###########
      self.disable_gradients(model_opt)
      self.enable_gradients(disc_opt)
  
      if self.mc_sampling:
          with torch.no_grad():
              disc_output = self.model(disc_images)
      else:
          disc_output = output
  
      fake_logits = self.discriminator(
          output.image.detach(), output.latent.detach()
      )
      real_logits = self.discriminator(
          disc_images, disc_output.latent.detach()
      )
  
      disc_loss = self.discriminator_loss(
          real_logits=real_logits,
          fake_logits=fake_logits,
          target=real_target,
      )
  
      # step optimizer
      disc_opt.zero_grad()
      self.manual_backward(disc_loss)
      if self.gradient_clip_val is not None:
          self.clip_gradients(
              disc_opt,
              gradient_clip_val=self.gradient_clip_val,
              gradient_clip_algorithm=self.gradient_clip_algorithm,
          )
      disc_opt.step()
      disc_sched.step()
  
      ########### QUANTILE OPTIMIZER ###########
      if quantile_opt is not None:
          quantile_loss = self.model.hyper_bottleneck.loss()
  
          quantile_opt.zero_grad()
          self.manual_backward(quantile_loss)
          quantile_opt.step()
          quantile_sched.step()
      else:
          quantile_loss = 0.0
  
      ########### LOGGING ###########
      self.log_quality_metrics(output, model_images, prefix="train")
      self.log_dict(
          {
              "train/total_model_loss": model_loss,
              "train/distortion_loss": distortion_loss,
              "train/rate_loss": rate_output.rate_loss,
              "train/generator_loss": g_loss,
              "train/discriminator_loss": disc_loss,
              "train/noisy_total_bpp": rate_output.total_bpp,
              "train/quantized_total_bpp": rate_output.quantized_total_bpp,
              "train/noisy_hyper_bpp": rate_output.hyper_bpp,
              "train/quantized_hyper_bpp": rate_output.quantized_hyper_bpp,
              "train/noisy_latent_bpp": rate_output.latent_bpp,
              "train/quantized_latent_bpp": rate_output.quantized_latent_bpp,
              "train/target": rate_output.target,
              "train/lam2": rate_output.lam2,
              "train/quantile_loss": quantile_loss,
          }
      )
  
imports: from typing import Optional, Tuple
  import hydra
  import torch
  import torch.nn
  from omegaconf import DictConfig
  from target_compression_module import TargetRateCompressionModule
  from torch import Tensor
  from neuralcompression import HyperpriorOutput
  from neuralcompression.loss_fn import DiscriminatorLoss, DistortionLoss, GeneratorLoss, TargetRateConfig
  from neuralcompression.models import HyperpriorAutoencoderBase
label_desc: Use image node_id 7 for calling the GANCompressionModule obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) without return types