Pytorch: How to assign a default value for look-up table using torch tensor
Say that I have two tensors as follows:

a = torch.tensor([[1, 2, 3], [1, 2, 3]])
b = torch.tensor([0, 2, 3, 4])


where b is the lookup value for a such as:

b[a]


will return the value of:

tensor([[2, 3, 4], [2, 3, 4]])


My problem is, what if I only have a look-up table of:

c = torch.tensor([0, 2, 3])


In which, for every out-of-index, I would like it to be assigned to index 0, such as c[a] will return

tensor([[2, 3, 0], [2, 3, 0]])


If I run c[a], of course, I will get this result:

RuntimeError: index 3 is out of bounds for dim with size 3


Thanks for your help.


  Code


# replace values greater than a certain number
def custom_replace(tensor, value, on_value):
    # we create a copy of the original tensor, 
    # because of the way we are replacing them.
    res = tensor.clone()
    res[tensor&gt;=value] = on_value
    return res

a = torch.tensor([[1, 2, 3], [1, 2, 3]])
c = torch.tensor([0, 2, 3])

a_ = custom_replace(a, c.size(0), 0)
print(c[a_])



  Output


tensor([[2, 3, 0],
        [2, 3, 0]])

