在 jupyter notebook 中如何指定使用指定 GPU 进行训练？

在 jupyter notebook 中如何指定使用指定 GPU 进行训练？

想在服务器上用训练好的模型测试效果，所以理想状态是直接使用 vscode 的 jupyter notebook 直接调用服务器的文件，免去下载到本地的麻烦。但是运行实验的时候由于无法指定 GPU 爆显存，故搜索在 jupyter notebook 中如何指定 GPU，但是貌似没人写。

也试过一些办法，比如：

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

但是都未能生效。

经过不懈搜索，真正起作用的是下面这条：

import torch
torch.cuda.set_device(3)  # 3 指的是第三块 GPU

注：我是在 vscode 上运行代码的，不知道浏览器中的 jupyter notebook 对其他方法能否生效。

推荐大家试一下 vscode 上的 jupyter notebook，省去了每次都要开命令行的烦恼。