code: def log_images(
      self, inputs: torch.Tensor, reconstructions: torch.Tensor
  ) -> Dict[str, torch.Tensor]:
      # calc logits of real/fake
      logits_real = self.discriminator(inputs.contiguous().detach())
      if len(logits_real.shape) < 4:
          # Non patch-discriminator
          return dict()
      logits_fake = self.discriminator(
          reconstructions.contiguous().detach()
      )
      # -> (b, 1, h, w)
  
      # parameters for colormapping
      high = max(
          logits_fake.abs().max(), logits_real.abs().max()
      ).item()
      cmap = colormaps["PiYG"]  # diverging colormap
  
      def to_colormap(logits: torch.Tensor) -> torch.Tensor:
          """(b, 1, ...) -> (b, 3, ...)"""
          logits = (logits + high) / (2 * high)
          logits_np = cmap(logits.cpu().numpy())[
              ..., :3
          ]  # truncate alpha channel
          # -> (b, 1, ..., 3)
          logits = torch.from_numpy(logits_np).to(logits.device)
          return rearrange(logits, "b 1 ... c -> b c ...")
  
      logits_real = torch.nn.functional.interpolate(
          logits_real,
          size=inputs.shape[-2:],
          mode="nearest",
          antialias=False,
      )
      logits_fake = torch.nn.functional.interpolate(
          logits_fake,
          size=reconstructions.shape[-2:],
          mode="nearest",
          antialias=False,
      )
  
      # alpha value of logits for overlay
      alpha_real = torch.abs(logits_real) / high
      alpha_fake = torch.abs(logits_fake) / high
      # -> (b, 1, h, w) in range [0, 0.5]
      # alpha value of lines don't really matter, since the values are the same
      # for both images and logits anyway
      grid_alpha_real = torchvision.utils.make_grid(alpha_real, nrow=4)
      grid_alpha_fake = torchvision.utils.make_grid(alpha_fake, nrow=4)
      grid_alpha = 0.8 * torch.cat(
          (grid_alpha_real, grid_alpha_fake), dim=1
      )
      # -> (1, h, w)
      # blend logits and images together
  
      # prepare logits for plotting
      logits_real = to_colormap(logits_real)
      logits_fake = to_colormap(logits_fake)
      # resize logits
      # -> (b, 3, h, w)
  
      # make some grids
      # add all logits to one plot
      logits_real = torchvision.utils.make_grid(logits_real, nrow=4)
      logits_fake = torchvision.utils.make_grid(logits_fake, nrow=4)
      # I just love how torchvision calls the number of columns `nrow`
      grid_logits = torch.cat((logits_real, logits_fake), dim=1)
      # -> (3, h, w)
  
      grid_images_real = torchvision.utils.make_grid(
          0.5 * inputs + 0.5, nrow=4
      )
      grid_images_fake = torchvision.utils.make_grid(
          0.5 * reconstructions + 0.5, nrow=4
      )
      grid_images = torch.cat(
          (grid_images_real, grid_images_fake), dim=1
      )
      # -> (3, h, w) in range [0, 1]
  
      grid_blend = (
          grid_alpha * grid_logits + (1 - grid_alpha) * grid_images
      )
  
      # Create labeled colorbar
      dpi = 100
      height = 128 / dpi
      width = grid_logits.shape[2] / dpi
      fig, ax = plt.subplots(figsize=(width, height), dpi=dpi)
      img = ax.imshow(np.array([[-high, high]]), cmap=cmap)
      plt.colorbar(
          img,
          cax=ax,
          orientation="horizontal",
          fraction=0.9,
          aspect=width / height,
          pad=0.0,
      )
      img.set_visible(False)
      fig.tight_layout()
      fig.canvas.draw()
      # manually convert figure to numpy
      cbar_np = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
      cbar_np = cbar_np.reshape(
          fig.canvas.get_width_height()[::-1] + (3,)
      )
      cbar = (
          torch.from_numpy(cbar_np.copy()).to(grid_logits.dtype) / 255.0
      )
      cbar = rearrange(cbar, "h w c -> c h w").to(grid_logits.device)
  
      # Add colorbar to plot
      annotated_grid = torch.cat((grid_logits, cbar), dim=1)
      blended_grid = torch.cat((grid_blend, cbar), dim=1)
      return {
          "vis_logits": 2 * annotated_grid[None, ...] - 1,
          "vis_logits_blended": 2 * blended_grid[None, ...] - 1,
      }
  
imports: from typing import Dict, Iterator, List, Optional, Tuple, Union
  import numpy
  import torch
  import torch.nn
  import torchvision
  from einops import rearrange
  from matplotlib import colormaps
  from matplotlib import pyplot
  from ....util import default, instantiate_from_config
  from ..lpips.loss.lpips import LPIPS
  from ..lpips.model.model import weights_init
  from ..lpips.vqperceptual import hinge_d_loss, vanilla_d_loss
label_desc: Use image node_id 4 for calling the GeneralLPIPSWithDiscriminator obj's underlying member method code with example usage: obj.log_images(inputs, reconstructions) and returns: dict, dict, rearrange