Solved by ptrblck in post #5 
              
              
                Yes your approach is right and I still think that the zero tensor is a valid answer, if your model just overfits to class0. 
Have a look at this example: 
x = torch.cat((torch.ones(1, 1, 5, 5), torch.zeros(1, 1, 5, 5)), 1)
print(x)
> tensor([[[[1., 1., 1., 1., 1.],
            [1., 1., 1., 1., 1.],
…