This should work:
x = torch.tensor([[[[ 0.4008, -0.6662],
                    [-0.4133,  1.3639]],

                   [[-0.8354,  0.6317],
                    [ 0.3240, -1.1438]],

                   [[-0.3452,  1.2110],
                    [ 0.6575,  0.9924]]]])

pred = torch.argmax(x, dim=1)
print(pred)
> tensor([[[0, 2],
           [2, 0]]])

out = torch.zeros_like(x).scatter_(1, pred.unsqueeze(1), 1.)
print(out)
> tensor([[[[1., 0.],
            [0., 1.]],

           [[0., 0.],
            [0., 0.]],

           [[0., 1.],
            [1., 0.]]]])