The fact that mask is only two dimensional is an issue if it should have color information. Can you try something like
mask = torch.zeros(3, h, w, dtype = torch.long)

and then
mask[0,:,:][validx] = torch.tensor(k[0], dtype=torch.long)
mask[1,:,:][validx] = torch.tensor(k[1], dtype=torch.long)
mask[2,:,:][validx] = torch.tensor(k[2], dtype=torch.long)

There might be a way to clean up the indexing here but I think this illustrates the intent.
The issue is that mapping is from colors to classes, but classes are being assigned to the mask here rather than colors. So here we are using the “key” of the mapping for the mask as that actually contains the colors.
On the other hand, if mask is supposed to have a format like 1000x1000 or 1x1000x1000 where each position is supposed to store the class id, then you might want to reuse part of the code here
  
    
    
    How to save multi-class segmentation prediction as image? vision
  
  
    Ah, curr_color should be used instead: 
    class_to_color = [torch.tensor([1.0, 0.0, 0.0]), ...]
    output = torch.zeros(1, 3, out.size(-2), out.size(-1), dtype=torch.float)
    for class_idx, color in enumerate(class_to_color):
        mask = out[:,class_idx,:,:] == torch.max(out, dim=1)[0]
        mask = mask.unsqueeze(1) # should have shape 1, 1, 100, 100
        curr_color = color.reshape(1, 3, 1, 1)
        segment = mask*curr_color # should have shape 1, 3, 100, 100
        output += seg…
  


to generate something that has 3 colors and is viewable