import torch


def convert_label_to_color(label, color_map):
  """Convert integer label to RGB image.
  """
  n, h, w = label.shape
  rgb = torch.index_select(color_map, 0, label.view(-1)).view(n, h, w, 3)
  rgb = rgb.permute(0, 3, 1, 2)

  return rgb