# 使用GPU或者CPU
#syntax 1
device = torch.device("cuda:0,1") 

#syntax 2  会根据设备情况自动选择cpu或者gpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 将model或者数据放入GPU
model.to(device)
model.gpu()
model.cuda()

# 将数据放入gpu
data = data.to(device)
data = data.gpu()
data = data.cuda()


# 指定GPU
# 终端指定
CUDA_VISIBLE_DEVICES=1 python script.py

# 在shell中指定
gpu=$1
echo "export CUDA_VISIBLE_DEVICES=${gpu}"
export CUDA_VISIBLE_DEVICES=${gpu}
python main.py  # 你要运行的python进程
执行sh run.sh 1,2即可使用1，2两张显卡。


# python 代码中指定
# 需放在代码开头部分
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2"


# 将model放在多个GPU上
model = nn.DataParallel(model)

参考链接
参考链接