torch.argwhere

torch.argwhere(input) → Tensor

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

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

Примечание

Эта функция похожа на NumPy’s argwhere.

Когда 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]])