请问应该在哪里设置显卡

[MengWH]

貌似显卡这边会报错 请问应该在哪里设置具体用哪张显卡

[xiaoye-hhh]

您好，我的机器只有一张卡就没有写。如果有多张卡，可以利用下面代码指定空闲的卡：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0" 

我这边没有多卡测试，如果没问题的话欢迎提交pr。

[MengWH]

非常感谢您的工作，目前问题已经解决。我上面报错的情况貌似是因为下面这段代码，torch.float16改成True就不报错了：