PyTorch torch.where 関数
Pytorch torch リファレンスマニュアルtorch.wherePyTorch において、条件に基づいて要素を返すために使用される関数です。
関数定義
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)
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]])
その他の拡張