python - 查找两个火炬张量的不相交

预先感谢大家的帮助！我想在PyTorch中做的是numpy的setdiff1d。例如，给出以下两个张量：

t1 = torch.tensor([1, 9, 12, 5, 24]).to('cuda:0')
t2 = torch.tensor([1, 24]).to('cuda:0')

预期输出应为（已排序或未排序）：

torch.tensor([9, 12, 5])

理想情况下，操作是在GPU上完成的，并且不会在GPU和CPU之间来回移动。非常感激！

最佳答案

如果您不想离开cuda，可以采用以下解决方法：

t1 = torch.tensor([1, 9, 12, 5, 24], device = 'cuda')
t2 = torch.tensor([1, 24], device = 'cuda')
indices = torch.ones_like(t1, dtype = torch.uint8, device = 'cuda')
for elem in t2:
    indices = indices & (t1 != elem)
intersection = t1[indices]