On the line
converted_tensor[i, sample == j] = x[j]
you just forgot to index on the channel dimension:
converted_tensor[i, :, sample == j] = x[j]
but then converted_tensor[i, :, mask] would be shaped [3, ?] whereas x[j] is shaped [3], so you would probably need to unsqueeze it:
converted_tensor[i, :, sample == j] = x[j].unsqueeze(1)

EDIT:
A much better way to do this would be to use torch.nn.functional.embedding() 6.
Put your color table in a [4, 3] float tensor:
x = torch.FloatTensor([[0, 0, 0], [255, 0, 0], [0, 0, 255], [0, 255, 0]])

Then you can directly have your result with:
converted_tensor = torch.nn.functional.embedding(network_output, x).permute(0, 3, 1, 2)

No need for a loop.