search

pytorch / extension-cpp

C++ extensions in PyTorch
1.02k stars 214 forks source link

Custom CUDA operator only work well on cuda:0 #78

Closed YuxianMeng closed 2 years ago

YuxianMeng commented 2 years ago

Hi! I'm learning how to write custom cuda operator by your tutorial.

However, I found the operator only output correct results on "cuda:0", and output all_zeros tensor in other devices like "cuda:1". What I did is just change this line to set device to "cuda:1")

iamgroot42 commented 2 years ago

I adapted the tutorial for my own code and ran into the same issue- did you find a fix for it?

iamgroot42 commented 2 years ago

@YuxianMeng I found a solution that works (it was buried inside closed issues) https://discuss.pytorch.org/t/c-cuda-extension-with-multiple-gpus/91241/6

YuxianMeng commented 2 years ago

@iamgroot42 Sorry for the late reply. I solved this issue with this post: https://discuss.pytorch.org/t/custom-cuda-operator-only-work-well-on-cuda-0/150504/7