pytorch指定gpu运行

直接在终端中

指定编号为1(默认从0开始)的GPU： CUDA_VISIBLE_DEVICES=1 python xxx.py

在脚本中指定

os.environ['CUDA_VISIBLE_DEVICES'] = gpus #gpus可以以数字的形式或者列表的形式

在脚本中操作

cuda = torch.device('cuda') # 默认CUDA设备（比如上一条的脚本指定gpu，或者在torch.cuda.device()里面等等情况） cuda0 = torch.device('cuda:0') cuda2 = torch.device('cuda:2') # GPU 2 (these are 0-indexed) #通过tensor.to()和tensor.cuda()的方式把tensor转移到特定设备 a = torch.randn(2).to(cuda2) #把tensor生成在GPU2上面 b = torch.randn(2).cuda(cuda2)