Solved by Mazhar_Shaikh in post #2 
              
              
                You can use this code 
color_map = #Tensor of shape(256,3)
gray_image = (gray_image * 255).long() # Tensor values between 0 and 255 and LongTensor and shape of (512,512)
output = color_map[gray_image] #Tensor of shape (512,512,3)

Why does this work? The gray_image tensor is used as an index in the …