Pytorch: Как назначить значение по умолчанию для справочной таблицы, используя тензор факела

Question

Скажите, что у меня есть два тензора следующим образом:

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

, где b - это значение поиска для a , например:

b[a]

вернет значение:

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

Моя проблема в том, что если у меня есть только справочная таблица:

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

В котором я хотел бы, чтобы для каждого вне индекса ему был присвоен индекс 0, например, c [a] вернет

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

Если я запускаю c [a] , конечно, я получу такой результат:

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

Спасибо за вашу помощь.

Answer 1

Код

# 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>=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_])

Выход

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