Thanks everyone in advance for your help! What I’m trying to do in PyTorch is something like numpy’s setdiff1d
. For example given the below two tensors:
JavaScript
x
3
1
t1 = torch.tensor([1, 9, 12, 5, 24]).to('cuda:0')
2
t2 = torch.tensor([1, 24]).to('cuda:0')
3
The expected output should be (sorted or unsorted):
JavaScript
1
2
1
torch.tensor([9, 12, 5])
2
Ideally the operations are done on GPU and no back and forth between GPU and CPU. Much appreciated!
Advertisement
Answer
if you don’t want to leave cuda, a workaround could be:
JavaScript
1
7
1
t1 = torch.tensor([1, 9, 12, 5, 24], device = 'cuda')
2
t2 = torch.tensor([1, 24], device = 'cuda')
3
indices = torch.ones_like(t1, dtype = torch.uint8, device = 'cuda')
4
for elem in t2:
5
indices = indices & (t1 != elem)
6
intersection = t1[indices]
7