dkoutsou:

If I use numpy then I won’t be able to run my code on GPU


You’d just have to temporarily copy tensor(s) to cpu, something like:
npx = x.detach().cpu().numpy()
npy = numpy_func(npx)
y = torch.from_numpy(npy).to(device=x.device)

overhead of this may be acceptable, YMMV



 dkoutsou:

Also intersect1d is not exactly what I want.


Check other set/sorting functions there.