torch.argwhere

torch.argwhere(input) → Tensor

Возвращает тензор, содержащий индексы всех ненулевых элементов input. Каждая строка результата содержит индексы ненулевого элемента input. Результат отсортирован лексикографически, причем последний индекс изменяется быстрее всего (в стиле C).

Если input имеет $n$ измерений, результирующий тензор индексов out имеет размер $(z \times n)$, где $z$ — общее количество ненулевых элементов в тензоре input.

Примечание

Эта функция аналогична argwhere в NumPy.

Если input находится на CUDA, эта функция вызывает синхронизацию между хостом и устройством.

Параметры:

{input} –

Пример:

>>> t = torch.tensor([1, 0, 1])
>>> torch.argwhere(t)
tensor([[0],
        [2]])
>>> t = torch.tensor([[1, 0, 1], [0, 1, 1]])
>>> torch.argwhere(t)
tensor([[0, 0],
        [0, 2],
        [1, 1],
        [1, 2]])