How to specify GPU usage?

[zhoubinxyz (Bin)]

Thanks a lot, it works
Hope pytorch can integrate with this argument to specify gpu usage.