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 0th dimension of color_map. The color_map tensor is accessed with this index and the new tensor created has the shape of gray_image inserted in the 0th dimension of output.