code: def apply_colormap(
      input_tensor: Tensor, colormap: ColorMap
  ) -> Tensor:
      r"""Apply to a gray tensor a colormap.
  
          The image data is assumed to be integer values.
  
      .. image:: _static/img/apply_colormap.png
  
      Args:
          input_tensor: the input tensor of a gray image.
          colormap: the colormap desired to be applied to the input tensor.
  
      Returns:
          A RGB tensor with the applied color map into the input_tensor
  
      Example:
          >>> input_tensor = torch.tensor([[[0, 1, 2], [25, 50, 63]]])
          >>> colormap = AUTUMN()
          >>> apply_colormap(input_tensor, colormap)
          tensor([[[1.0000, 1.0000, 1.0000],
                   [1.0000, 1.0000, 1.0000]],
          <BLANKLINE>
                  [[0.0000, 0.0159, 0.0317],
                   [0.3968, 0.7937, 1.0000]],
          <BLANKLINE>
                  [[0.0000, 0.0000, 0.0000],
                   [0.0000, 0.0000, 0.0000]]])
      """
      # FIXME: implement to work with RGB images
      # should work with KORNIA_CHECK_SHAPE(x, ["B","C", "H", "W"])
  
      KORNIA_CHECK_IS_GRAY(input_tensor)
  
      if (
          len(input_tensor.shape) == 4 and input_tensor.shape[1] == 1
      ):  # if (B x 1 X H x W)
          input_tensor = input_tensor[:, 0, ...]  # (B x H x W)
      elif (
          len(input_tensor.shape) == 3 and input_tensor.shape[0] == 1
      ):  # if (1 X H x W)
          input_tensor = input_tensor[0, ...]  # (H x W)
  
      keys = torch.arange(
          0,
          len(colormap) - 1,
          dtype=input_tensor.dtype,
          device=input_tensor.device,
      )  # (num_colors)
  
      index = torch.bucketize(
          input_tensor, keys
      )  # shape equals <input_tensor>: (B x H x W) or (H x W)
  
      output = colormap.colors[
          :, index
      ]  # (3 x B x H x W) or (3 x H x W)
  
      if len(output.shape) == 4:
          output = output.permute(1, 0, -2, -1)  # (B x 3 x H x W)
  
      return output
  
file_path: kornia/kornia/color/colormap.py
imports: from abc import ABC
  from typing import List, Optional
  import torch
  from torch.nn.functional import interpolate
  from kornia.core import Module, Tensor, tensor
  from kornia.core.check import KORNIA_CHECK_IS_GRAY
label_desc: Use image node_id 2 for calling a global function with example usage: apply_colormap(input_tensor, colormap) and returns: output