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

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

img

问题

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

可能需要在代码里指定吧

你把项目里面的--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')