Hi Yun and @ptrblck!



 ptrblck:

out = torch.zeros_like(x).scatter_(1, pred.unsqueeze(1), 1.)



As an alternative, you could use torch.nn.functional.one_hot(),
but would now have to permute() the dimensions of the result to get
them back in the order you want them:
>>> import torch
>>> torch.__version__
'1.9.0'
>>>
>>> t = 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]]]
... ])
>>>
>>> ind = t.argmax (dim = 1)
>>>
>>> ind
tensor([[[0, 2],
         [2, 0]]])
>>>
>>> torch.nn.functional.one_hot (ind).permute (0, 3, 1, 2)
tensor([[[[1, 0],
          [0, 1]],

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

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

This is arguably a bit more readable, as with “one_hot()” you say
what you mean (at the cost of the permute()).
(Note that, as @ptrblck illustrated, you do not need to use softmax().
Because it doesn’t change the order of the values, leaving it out
doesn’t change the result of argmax().)
Best.
K. Frank