Thanks @ptrblck. One more question (sorry for troubling. I am working with multi-class problem for the first time)
I am looking at the output:

        X = normalize_zero_to_one(X)
        y = normalize_zero_to_one(y)
        
        images = Variable(torch.from_numpy(X)).to(device) # [batch, channel, H, W]
        masks = Variable(torch.from_numpy(y)).to(device)
        
        masks = torch.argmax(masks, dim=1)
        outputs = model(images)            
        loss = loss_new(outputs, masks)          
        
        preds = torch.argmax(outputs, dim=1)
        
        print(preds.shape)
        print(preds)



torch.Size([1, 256, 256])
tensor([[[1, 1, 1,  …, 1, 1, 1],
[1, 1, 1,  …, 1, 1, 1],
[1, 1, 1,  …, 1, 1, 0],
…,
[1, 1, 1,  …, 1, 1, 1],
[1, 1, 1,  …, 1, 1, 0],
[1, 0, 0,  …, 1, 1, 0]]])

Can you help me understand how to read this output? When I see the first value as 1, how do I read which output band is it talking about
P.S. Note, for certain simplicity, I changed my output classes to 2 classes (  Band5','Band6) and got output shape [batch_size, nb_classes=2, height, width]. Print result in the answer reflects the same.