Javier:

What I would like to do is to binarize along the x dimension (dim=2) for any batch or channel, i.e for every yz location I want to set the maximum value along x-axis to 1 and the rest to zero.


If I understand this use case correctly, this should work:
x = torch.randn(2, 3, 4, 4, 4)
idx = torch.argmax(x, dim=2, keepdims=True)
ret = torch.zeros_like(x).scatter_(2, idx, 1.)