I cannot reproduce that issue.
Can you check the shape of your output is expected?
>>> out = torch.randn(1, 6, 100, 100)
>>> mask = out[:,0,:,:] == torch.max(out, dim=1)[0]
>>> mask.shape
torch.Size([1, 100, 100])
>>> mask.unsqueeze(1).shape
torch.Size([1, 1, 100, 100])