code: def __init__(self):
      super().__init__()
      self.custom_print_logger = logging.getLogger(
          self.__class__.__name__
      )
  
      # train metrics
      self.train_psnr = PeakSignalNoiseRatio(data_range=255.0)
      self.train_mse = MeanSquaredError()
      self.train_msssim = MultiscaleStructuralSimilarity(
          data_range=255.0
      )
  
      # validation metrics
      self.eval_psnr = PeakSignalNoiseRatio(data_range=255.0)
      self.eval_mse = MeanSquaredError()
      self.eval_msssim = MultiscaleStructuralSimilarity(
          data_range=255.0
      )
      self.eval_lpips = LearnedPerceptualImagePatchSimilarity(
          normalize=True
      )
      self.eval_fid = FrechetInceptionDistance()
  
      # test metrics, set these to None to start
      self.test_metrics_set = False
      self.test_mse: Optional[MeanSquaredError] = None
      self.test_msssim: Optional[MultiscaleStructuralSimilarity] = None
      self.test_lpips: Optional[
          LearnedPerceptualImagePatchSimilarity
      ] = None
      self.test_dists: Optional[
          ncm.DeepImageStructureTextureSimilarity
      ] = None
      self.test_fid: Optional[ncm.FrechetInceptionDistance] = None
      self.test_fid_swav: Optional[
          ncm.FrechetInceptionDistanceSwAV
      ] = None
      self.test_kid: Optional[ncm.KernelInceptionDistance] = None
  
      self.hific_mse_param = 0.075 * 2**-5
      self.hific_lpips_param = 1.0
      self.test_dataset_name: Optional[str] = None
      self.patch_size = 256
  
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 1 to create a new ImageModule object from inherited base classes: LightningModule with example: obj = ImageModule()