Ive trained my segmentation network using a U-Net model. I want to now visualize the result from the network output.
Theshape of output is [1,2,256,256]. One channel for background and the other for foreground. Here is what I get for print(output.detach().squeeze()):
tensor([[[ 6.1109,  6.1109,  6.1126,  ...,  6.1580,  6.1317,  6.1317],
         [ 6.1109,  6.1109,  6.1126,  ...,  6.1580,  6.1317,  6.1317],
         [ 6.1125,  6.1125,  6.1139,  ...,  6.1281,  6.1075,  6.1075],
         ...,
         [ 8.3482,  8.3482,  8.3518,  ...,  8.7196,  8.7172,  8.7172],
         [ 8.3446,  8.3446,  8.3485,  ...,  8.7216,  8.7194,  8.7194],
         [ 8.3446,  8.3446,  8.3485,  ...,  8.7216,  8.7194,  8.7194]],

        [[-6.8949, -6.8949, -6.8821,  ..., -7.1064, -7.0877, -7.0877],
         [-6.8949, -6.8949, -6.8821,  ..., -7.1064, -7.0877, -7.0877],
         [-6.9068, -6.9068, -6.8939,  ..., -7.0928, -7.0780, -7.0780],
         ...,
         [-8.3679, -8.3679, -8.3719,  ..., -8.7457, -8.7435, -8.7435],
         [-8.3656, -8.3656, -8.3700,  ..., -8.7487, -8.7467, -8.7467],
         [-8.3656, -8.3656, -8.3700,  ..., -8.7487, -8.7467, -8.7467]]])

To get the predictions from this output I did:
predicted = torch.argmax(output.detach(), 1).squeeze(0)
But this gives a tensor with all zeros??
I then write this image along with the respective label using:
f_ = str(idx) + '_predicted.jpg'
io.imwrite(os.path.join(output_folder, f_), predicted.float().cpu().numpy())

f_ = str(idx) + '_target.jpg'
io.imwrite(os.path.join(output_folder, f_), labels.squeeze().float().cpu().numpy())

The shape of predicted and label(after squeeze())  is [256 256].
Label is being written fine as can be seen:

But I am only getting black for predicted:

Am I doing this right? How else am I supposed to visualize the network output?
Thanks