code: def add_heatmap(tensor):
      """
      Add heatmap to 2D tensor.
      Args:
          tensor (tensor): a 2D tensor. Tensor value must be in [0..1] range.
      Returns:
          heatmap (tensor): a 3D tensor. Result of applying heatmap to the 2D tensor.
      """
      assert tensor.ndim == 2, "Only support 2D tensors."
      # Move tensor to cpu if necessary.
      if tensor.device != torch.device("cpu"):
          arr = tensor.cpu()
      else:
          arr = tensor
      arr = arr.numpy()
      # Get the color map by name.
      cm = plt.get_cmap("viridis")
      heatmap = cm(arr)
      heatmap = heatmap[:, :, :3]
      # Convert (H, W, C) to (C, H, W)
      heatmap = torch.Tensor(heatmap).permute(2, 0, 1)
      return heatmap
  
imports: import logging
  import math
  import os
  import matplotlib.pyplot
  import torch
  from torch.utils.tensorboard import SummaryWriter
  from torchvision.utils import make_grid
  import slowfast.utils.logging
  import slowfast.visualization.utils
  from slowfast.utils.misc import get_class_names
label_desc: Use image node_id 4 for calling a global function with example usage: add_heatmap(tensor) and returns: heatmap