search

fxia22 / stn.pytorch

pytorch version of spatial transformer networks
Other
589 stars 87 forks source link

Bug in gpu support #9

Open maga33 opened 7 years ago

maga33 commented 7 years ago

Hi, Thank you for sharing great code. I found that the code works fine with gpu_id = 0, but it raises cuda runtime error if I tried to run it on different gpus. I guess some parts of codes are implicitly assuming to use default gpu, which would be great to be modified in the future.

fxia22 commented 7 years ago

I just updated it to support non-default device id. An example can be found in test.py. Let me know if that works for you, thanks a lot!