Функция torch.where в PyTorch

Pytorch torch 参考手册Справочное руководство Pytorch torch

torch.whereЭто функция в PyTorch, используемая для возврата элементов на основе условий.

Определение функции

torch.where(condition, input, other)

Пример использования

Пример

import torch

condition = torch.tensor([[True, False], [False, True]])
x = torch.tensor([[1, 2], [3, 4]])
y = torch.tensor([[10, 20], [30, 40]])

result = torch.where(condition, x, y)

print(result)

Результат вывода:

tensor([[ 1, 20],
        [30,  4]])

Pytorch torch 参考手册Справочное руководство Pytorch torch

Другие расширения