（转载）Pytorch运行时指定对应序号的显卡

Pytorch运行时指定对应序号的显卡

原文链接：https://zhuanlan.zhihu.com/p/166161217

1. 在代码中指定

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

2. 在命令行中执行代码时指定

1
CUDA_VISIBLE_DEVICES='2' python3 train.py

3. 在命令行中执行脚本文件时指定

1
CUDA_VISIBLE_DEVICES='2' sh run.sh