使用torch使用GPU时，报错RuntimeError: device >= 0 && device < num_gpus INTERNAL...

使用torch使用GPU时，报错RuntimeError: device >= 0 && device < num_gpus INTERNAL...
2023-09-17
639
阅读1分钟

错误原因： torch在设置

os.environ['CUDA_VISIBLE_DEVICES']
就被import了。这样torch实际上就看到了所有的卡。

解决方案：

import os

这个一定要在import torch之前！否则torch会能看到所有的gpu。

os.environ['CUDA_VISIBLE_DEVICES'] = '3,4'