How to set gpu id in train.py

genforce / idinvert

[ECCV 2020] In-Domain GAN Inversion for Real Image Editing

https://genforce.github.io/idinvert/

MIT License

461 stars 65 forks source link

How to set gpu id in train.py #39

Closed hadzica closed 3 years ago

hadzica commented 3 years ago

Hi guys

Thanks for providing your great work. I was able to run train.py. But I'm wondering how to set the gpu id in train.py. Do you have a suggestion. I would be very thankful.

Thanks

ShenYujun commented 3 years ago

Please try os.environ['CUDA_VISIBLE_DEVICES'] = '0'

hadzica commented 3 years ago

thanks!