How can I get argmaxed torch tensor excluding certain index?

user_10919971 0
Viewed 729

I wonder if I can get torch.argmax of my input excluding certain index. For example,

target = torch.tensor([1,2])
input = torch.tensor([[0.1,0.5,0.2,0.2], [0.1,0.5,0.1,0.3]])

I want to get the maximum value in input excluding the index on the target, so that the result would be

output = torch.tensor([[0.2],[0.5]])
pytorch
2 Answers

You can try this

  • Set negative infy to the target indices in temp tensor
  • Then use torch.max or torch.argmax
tmp_input = input.clone()
tmp_input[range(len(input)), target] = float("-Inf")

torch.max(tmp_input, dim=1).values
tensor([0.2000, 0.5000])

torch.max(tmp_input, dim=1).indices
tensor([3, 1])

torch.argmax(tmp_input, dim=1)
tensor([3, 1])
user_60756990 
input[target[0]-1,target[1]-1] = -1 # or use -inf 
#-1 is added for python indexing style
output  = torch.max(input,dim = 1)
user_130730200
Related