Bug in gpu support

[maga33]

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]

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!