Setting the top 1 element of a tensor to be 1

1

Hi! Given a 1D tensor, I would need to find its biggest number and turn that index into value 1, and all the rest of the indices are all set to 0.

For example I have [1,2,4,6,2,1], and I would like to convert it into [0,0,0,1,0,0].

I’ve used
idx = torch.topk(tensor,1)[1]
out = torch.zeros_like(tensor)
out[idx] = 1

But I’m not sure if there are better ways.

Thank you!

2

Hi,

You can try this as well.

import torch

tensor = torch.tensor([1,2,4,6,2,1])
idx = torch.topk(tensor,1)[1]
out = (tensor >= tensor[idx]).type(torch.uint8)

Thanks

3

Yeah, much neater! Thanks!