code: def training_step(self, batch, batch_idx):
      images: Tensor
      output: HyperpriorOutput
      model_opt, quantile_opt = self.optimizers()
      model_sched, quantile_sched = self.lr_schedulers()
  
      ########### MODEL OPTIMIZER ###########
      images = batch
  
      output = self.model(images)
  
      distortion_loss = self.distortion_loss(output.image, images)
      rate_output = self.rate_loss(
          original=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,
      )
      loss_train = (
          self.distortion_lam * distortion_loss + rate_output.rate_loss
      )
  
      model_opt.zero_grad()
      self.manual_backward(loss_train)
      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()
  
      ########### QUANTILE OPTIMIZER ###########
      quantile_loss = self.model.hyper_bottleneck.loss()
  
      quantile_opt.zero_grad()
      self.manual_backward(quantile_loss)
      quantile_opt.step()
      quantile_sched.step()
  
      ########### LOGGING ###########
      self.log_quality_metrics(output, images, prefix="train")
      self.log_dict(
          {
              "train/total_loss": loss_train,
              "train/distortion_loss": distortion_loss,
              "train/rate_loss": rate_output.rate_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: import math
  from typing import List, NamedTuple, Optional
  import hydra
  import torch
  from image_module import ImageModule
  from omegaconf import DictConfig
  from torch import Tensor
  from torch.utils.data import Dataset
  from neuralcompression import HyperpriorOutput
  from neuralcompression.loss_fn import DistortionLoss, MSELoss, TargetRateConfig
  from neuralcompression.metrics import pickle_size_of
  from neuralcompression.models import HyperpriorAutoencoderBase
label_desc: Use image node_id 8 for calling the TargetRateCompressionModule obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) without return types