ptrblck:

target = torch.argmax(target, dim=1)


Ok! But currently my target shape is (1, 3, 256, 256), isn’t it expecting something like (1, 1, 256, 256) for me to apply target = torch.argmax(target, dim=1)?
I’ll try the code snippet in the meanwhile.