@ptrblck Thank you for the help! I see the loop works for the loss but how would we get the prediction values for each class from it?
Edit 1:
My current test code set up looks like this but how to get the right prediction:

        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)          
        output_sigmoid = torch.sigmoid(outputs)
        preds = output_sigmoid > output_sigmoid.cpu().numpy().mean()


I don’t think I’ll need sigmoid here, isn’t it?