I am trying to implement the conversion of gray scale image to an RGB image through a colormap myself in PyTorch. I have:

- A tensor
*gray_image*which is a LongTensor of 512 x 512 with (integer) values between 0 and 255 - The
*color_map*, which is a FloatTensor of 256 x 3, where every row represents the RGB value for a specific gray value - An empty output file, which is a FloatTensor of 512 x 512 x 3

I fill the output tensor like this:

```
output = torch.zeros((512, 512 3) # Initialize empty
gray_image = (gray_image * 255).long() # Tensor values between 0 and 255 and LongTensor
for i in range(gray_image.shape[0]):
for j in range(gray_image.shape[1]):
output[i, j, :] = color_map[grads[i, j]]
```

This is quite slow and this must be possible in 1 or 2 lines using `torch.gather()`

. I tried to do this myself, but did not succeed and found the documentation of `torch.gather()`

to be quite complex. Can someone help me with this?