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: [SOLVED] Index to rgb , tensor casting from given cmap , can close that.