code: def log_quality_metrics(
      self,
      output: Union[HyperpriorOutput, Tensor],
      input_images: Tensor,
      prefix,
      dataset_name: Optional[str] = None,
  ):
      # update metrics
      if isinstance(output, Tensor):
          pred = output.clamp(0.0, 1.0)
      else:
          assert output.image is not None
          pred = output.image.clamp(0.0, 1.0)
      rescaled_pred = ncF.image_to_255_scale(pred)
      rescaled_input_images = ncF.image_to_255_scale(input_images)
  
      batch_size = pred.shape[0]
  
      if prefix == "train":
          with torch.no_grad():
              self.train_psnr(rescaled_pred, rescaled_input_images)
              self.train_msssim(rescaled_pred, rescaled_input_images)
              self.train_mse(
                  rescaled_pred.view(batch_size, -1),
                  rescaled_input_images.view(batch_size, -1),
              )
          self.log(f"{prefix}/psnr", self.train_psnr, prog_bar=True)
          self.log(f"{prefix}/mse", self.train_mse, prog_bar=True)
          self.log(f"{prefix}/ms_ssim", self.train_msssim)
      elif prefix == "val":
          self.eval_psnr(rescaled_pred, rescaled_input_images)
          self.eval_msssim(rescaled_pred, rescaled_input_images)
          mse_val = self.eval_mse(
              rescaled_pred.view(batch_size, -1),
              rescaled_input_images.view(batch_size, -1),
          )
          lpips_val = self.eval_lpips(pred, input_images)
          hific_distortion = (
              self.hific_mse_param * mse_val
              + self.hific_lpips_param * lpips_val
          )
          self.eval_fid.update(
              ncF.image_to_255_scale(input_images, dtype=torch.uint8),
              real=True,
          )
          self.eval_fid.update(
              ncF.image_to_255_scale(pred, dtype=torch.uint8),
              real=False,
          )
          self.log_dict(
              {
                  f"{prefix}/psnr": self.eval_psnr,
                  f"{prefix}/mse": self.eval_mse,
                  f"{prefix}/ms_ssim": self.eval_msssim,
                  f"{prefix}/lpips": self.eval_lpips,
                  f"{prefix}/hific_distortion": hific_distortion,
              },
              sync_dist=True,
          )
      elif prefix == "test":
          assert self.test_msssim is not None
          assert self.test_mse is not None
          assert self.test_lpips is not None
          assert self.test_dists is not None
          if dataset_name is None:
              raise ValueError("Must pass dataset_name for test")
  
          psnr_val = ncm.calc_psnr(rescaled_pred, rescaled_input_images)
          self.test_msssim(rescaled_pred, rescaled_input_images)
          mse_val = self.test_mse(
              rescaled_pred.view(batch_size, -1),
              rescaled_input_images.view(batch_size, -1),
          )
          lpips_val = self.test_lpips(pred, input_images)
          self.test_dists(pred, input_images)
          hific_distortion = (
              self.hific_mse_param * mse_val
              + self.hific_lpips_param * lpips_val
          )
          clic_eval = "clic" in dataset_name or "div2k" in dataset_name
          self._update_test_fid_kid(
              input_images=input_images, pred=pred, clic_eval=clic_eval
          )
          self.log_dict(
              {
                  f"{prefix}/{dataset_name}_psnr": psnr_val,
                  f"{prefix}/{dataset_name}_mse": self.test_mse,
                  f"{prefix}/{dataset_name}_ms_ssim": self.test_msssim,
                  f"{prefix}/{dataset_name}_lpips": self.test_lpips,
                  f"{prefix}/{dataset_name}_dists": self.test_dists,
                  f"{prefix}/{dataset_name}_hific_distortion": hific_distortion,
              }
          )
      else:
          raise ValueError(f"Unrecognized prefix {prefix}")
  
imports: import logging
  from typing import Optional, Union
  import torch
  from lightning.pytorch import LightningModule
  from torch import Tensor
  from torchmetrics import MeanSquaredError, PeakSignalNoiseRatio
  from torchmetrics.image import FrechetInceptionDistance, LearnedPerceptualImagePatchSimilarity
  import neuralcompression.functional
  import neuralcompression.metrics
  from neuralcompression import HyperpriorOutput
  from neuralcompression.metrics import MultiscaleStructuralSimilarity
label_desc: Use image node_id 6 for calling the ImageModule obj's underlying member method code with example usage: obj.log_quality_metrics(output, input_images, prefix, dataset_name) without return types