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.],
            [1., 1., 1., 1., 1.],
            [1., 1., 1., 1., 1.],
            [1., 1., 1., 1., 1.]],

           [[0., 0., 0., 0., 0.],
            [0., 0., 0., 0., 0.],
            [0., 0., 0., 0., 0.],
            [0., 0., 0., 0., 0.],
            [0., 0., 0., 0., 0.]]]])
print(torch.argmax(x, 1))
> tensor([[[0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0]]])

A zero tensor is in this case the right answer, since all values in channel0 are larger than in channel1.
If we manipulate a certain value in channel1, you’ll also get the right answer:
x[0, 1, 0, 0] = 2
print(torch.argmax(x, 1))
> tensor([[[1, 0, 0, 0, 0],
           [0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0],
           [0, 0, 0, 0, 0]]])