PyTorch torch.isin 函数
torch.isin 是 PyTorch 中用于检查元素是否在另一个张量中的函数。返回一个布尔张量,表示第一个张量的每个元素是否出现在第二个张量中。
函数定义
torch.isin(elements, test_elements, assume_unique=False, invert=False)
使用示例
实例
import torch
# 创建两个张量
elements = torch.tensor([1, 2, 3, 4])
test_elements = torch.tensor([2, 4])
# 检查元素是否在test_elements中
result = torch.isin(elements, test_elements)
print(result)
# 创建两个张量
elements = torch.tensor([1, 2, 3, 4])
test_elements = torch.tensor([2, 4])
# 检查元素是否在test_elements中
result = torch.isin(elements, test_elements)
print(result)
输出结果为:
tensor([False, True, False, True])

Pytorch torch 参考手册