Hi @ptrblck
Even in that case I should get some output from argmax instead of a zero tensor. The way I am visualizing the outputs here was suggested by you in my previous question: Understanding how to label/target tensors for 3D volumes
And this seems to work fine for 3D volumes. But here I am using 2D images. But the concept should be the same and it should work. As I have labeled my label tensor during training such that channel0 represents background and channel1 represents my foreground.
So shouldnt I be looking for the maximum values in dim=1 ?