def apply_colormap(input_tensor: Tensor, colormap: ColorMap) -> Tensor:
    r"""Apply to a gray tensor a colormap.

    .. image:: _static/img/apply_colormap.png

    Args:
        input_tensor: the input tensor of 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.

    Raises:
        ValueError: If `colormap` is not a ColorMap object.

    .. note::
        The input tensor must be integer values in the range of [0-255] or float values in the range of [0-1].

    Example:
        >>> input_tensor = torch.tensor([[[0, 1, 2], [15, 25, 33], [128, 158, 188]]])
        >>> colormap = ColorMap(base=ColorMapType.autumn)
        >>> apply_colormap(input_tensor, colormap)
        tensor([[[[1.0000, 1.0000, 1.0000],
                  [1.0000, 1.0000, 1.0000],
                  [1.0000, 1.0000, 1.0000]],
        <BLANKLINE>
                 [[0.0000, 0.0159, 0.0159],
                  [0.0635, 0.1111, 0.1429],
                  [0.5079, 0.6190, 0.7302]],
        <BLANKLINE>
                 [[0.0000, 0.0000, 0.0000],
                  [0.0000, 0.0000, 0.0000],
                  [0.0000, 0.0000, 0.0000]]]])

    """
    KORNIA_CHECK(isinstance(input_tensor, Tensor), f"`input_tensor` must be a Tensor. Got: {type(input_tensor)}")
    valid_types = [torch.half, torch.float, torch.double, torch.uint8, torch.int, torch.long, torch.short]
    KORNIA_CHECK(
        input_tensor.dtype in valid_types, f"`input_tensor` must be a {valid_types}. Got: {input_tensor.dtype}"
    )
    KORNIA_CHECK(len(input_tensor.shape) in (3, 4), "Wrong input tensor dimension.")
    if len(input_tensor.shape) == 3:
        input_tensor = input_tensor.unsqueeze_(0)

    B, C, H, W = input_tensor.shape
    input_tensor = input_tensor.reshape(B, C, -1)
    max_value = 1.0 if input_tensor.max() <= 1.0 else 255.0
    input_tensor = input_tensor.float().div_(max_value)

    colors = colormap.colors.permute(1, 0)
    num_colors, channels_cmap = colors.shape
    keys = torch.linspace(0.0, 1.0, num_colors - 1, device=input_tensor.device, dtype=input_tensor.dtype)
    indices = torch.bucketize(input_tensor, keys).unsqueeze(-1).expand(-1, -1, -1, 3)

    output = torch.gather(colors.expand(B, C, -1, -1), 2, indices)
    # (B, C, H*W, channels_cmap) -> (B, C*channels_cmap, H, W)
    output = output.permute(0, 1, 3, 2).reshape(B, C * channels_cmap, H, W)

    return output
Apply to a gray tensor a colormap.

    .. image:: _static/img/apply_colormap.png

    Args:
        input_tensor: the input tensor of 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.

    Raises:
        ValueError: If `colormap` is not a ColorMap object.

    .. note::
        The input tensor must be integer values in the range of [0-255] or float values in the range of [0-1].

    Example:
        >>> input_tensor = torch.tensor([[[0, 1, 2], [15, 25, 33], [128, 158, 188]]])
        >>> colormap = ColorMap(base=ColorMapType.autumn)
        >>> apply_colormap(input_tensor, colormap)
        tensor([[[[1.0000, 1.0000, 1.0000],
                  [1.0000, 1.0000, 1.0000],
                  [1.0000, 1.0000, 1.0000]],
        <BLANKLINE>
                 [[0.0000, 0.0159, 0.0159],
                  [0.0635, 0.1111, 0.1429],
                  [0.5079, 0.6190, 0.7302]],
        <BLANKLINE>
                 [[0.0000, 0.0000, 0.0000],
                  [0.0000, 0.0000, 0.0000],
                  [0.0000, 0.0000, 0.0000]]]])

    
kornia/color/colormap.py