This should do the trick
import torch
a = torch.Tensor([[0.0, 0.9, 0.8, 0.0, 0.2, 0.0],
                  [0.2, 0.0, 0.0, 0.8, 0.7, 0.0],
                  [0.6, 0.0, 0.0, 0.0, 0.0, 0.7]])
b = torch.argmax(a,dim=1).unsqueeze(1)
one_hot = torch.zeros(a.shape).scatter(1,b,1)
req_op = a * one_hot