ptrblck:

pred_imgs = [cityscapes_map[p] for p in pred]


Thanks a lot, that works perfectly.
though I used : _, pred = torch.max(output, dim=1)  from one of your previous answers, does it have a difference from using argmax ?
And you also gave me an answer for this: Index to rgb , tensor casting from given cmap 57 , can close that.