I have a tensor:
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]])

and I want to extract the max value of dim=1, and keep it’s original position, like this:
res = torch.Tensor([[0.0, 0.9, 0.0, 0.0, 0.0, 0.0],
                  [0.0, 0.0, 0.0, 0.8, 0.0, 0.0],
                  [0.0, 0.0, 0.0, 0.0, 0.0, 0.7]])

How to implement this?
Thank you!