Solved by ptrblck in post #4 
              
              
                This code snippet should work: 
x = torch.tensor([[2.3, 3.1, 1.2],
                  [5.6, 4.1, 3.1]])
idx = torch.argmax(x, 1)
output = torch.zeros_like(x).scatter_(1, idx[:, None], 1.)