stn.pytorch

stn.pytorch copied to clipboard

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.

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!