怎么设置参数指定gpu运行pytorch_lightning程序

Question

问题遇到的现象和发生背景

问题

上面是要运行pytorch_lightning写的main函数，但需要传入使用的gpu设备信息，没怎么接触过pytorch_lightning，如果我想指定某一个gpu执行，比如说cuda0设备运行上述代码，那个gpus出的代码应该怎么修改？

Answer 1

你把项目里面的--gpus 1改成--gpus 0就可以调用cuda0了。
如果你进到main.py里面去，搜索下这个gpus，应该可以看到一个类似的参数,如果里面没有default，加上一个默认参数，就可以不加--gpus 0就可以默认cuda0了

parser.add_argument('--gpus', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')