今天在进行模型训练的时候遇到这样一个问题:我想要模型在第1块显卡上进行训练,但是不论怎么指定GPU,模型始终是在第0块显卡上进行训练,很是不理解为什么这样?
经过查找资料,发现一个问题:如果在torch之后再指定GPU,那么这个指定是无效的,究竟为什么我还是没有想通,如果有知道原因的小伙伴麻烦下面评论告诉我一下,谢谢!
正确代码如图所示:
import os os.environ["CUDA_VISIBLE_DEVICES"] = '1' import torch