在pytorch中，我如何将第二个gpu设置为juptyer笔记本中的默认gpu

Question

我有两个gpu，想要打开第二个juptyer笔记本，并确保其中的所有内容都只在第二个gpu上运行，而不是在第一个上运行。理想情况下，我希望在开始时运行一个单元格，而不是在多个位置传递device=1。

Answer 1

在这种情况下，答案是调用set设备。

import torch
torch.cuda.set_device(1) 

然而，文档并不鼓励这种做法(https://github.com/pytorch/pytorch/issues/260)我的场景是唯一合理的用例。首选方法是设置环境变量，这在笔记本电脑的情况下也是可能的：

import os
os.environ['CUDA_VISIBLE_DEVICES']='1'