code: def test_step(self, batch, batch_idx, dataloader_idx=0):
      prefix = "test"
      images: Tensor
      if self.test_dataset_name is None:
          raise ValueError("Must run module.update_test_dataset_name()")
  
      self.model.update_tensor_devices("compress")
  
      dataset_name = self.test_dataset_name
  
      if isinstance(batch, Tensor):
          images = batch
      else:
          images, _ = batch
  
      compressed = self.model.compress(images, force_cpu=False)
      decompressed = self.model.decompress(compressed, force_cpu=False)
  
      test_loss = self.distortion_loss(decompressed, images)
  
      num_bytes = pickle_size_of(compressed)
      bpp = (
          num_bytes
          * 8
          / (images.shape[0] * images.shape[-2] * images.shape[-1])
      )
  
      self.log_dict(
          {
              f"{prefix}/{dataset_name}_bpp": bpp,
              f"{prefix}/{dataset_name}_test_distortion": test_loss,
          }
      )
      self.log_quality_metrics(
          decompressed, images, prefix=prefix, dataset_name=dataset_name
      )
  
      return test_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 12 for calling the TargetRateCompressionModule obj's underlying member method code with example usage: obj.test_step(batch, batch_idx, dataloader_idx) and returns: test_loss