指定GPU设备

指定GPU设备

EchoL、2026-01-10 10:32

在运行代码的时候遇到下面报错：

完整报错信息：

python 复制代码

DeferredCudaCallError: CUDA call failed lazily at initialization with error: device >= 0 && device < num_gpus INTERNAL ASSERT FAILED at "/opt/conda/conda-bld/pytorch_1716905969824/work/aten/src/ATen/cuda/CUDAContext.cpp":50, please report a bug to PyTorch. device=, num_gpus=

排查后发现问题出在有1块GPU error了，在代码中显式地指定GPU编号即可解决问题。

方法一：

python 复制代码

import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2'

CUDA_VISIBLE_DEVICES ：决定进程能看见哪些物理 GPU，并且会重新编号（映射成从 0 开始的逻辑编号），要放在import torch前。

方法二：

python 复制代码

import torch
torch.cuda.set_device(2)

torch.cuda.set_device(k) 是在"已经可见的逻辑 GPU 列表"里，把默认 GPU 指针设成第 k 个；不会屏蔽 GPU，也不会改变可见性。

简单来说：

想"只用某几张卡 / 不要碰其他卡" → 用CUDA_VISIBLE_DEVICES来设置；
想"在可见卡里把默认卡设成哪张" → 用 set_device；

示例1：

python 复制代码

os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
import torch
torch.cuda.set_device(1)

这里其实是在选：逻辑1->GPU 3，而不是GPU 1。

示例2：

python 复制代码

os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
import torch
torch.cuda.set_device(2)   # 报错

因为可见逻辑设备只有0和1，没有2。

参考：

$1$ 【设置gpu设备】os.environ $'CUDA_VISIBLE_DEVICES'$ 和 torch.cuda.set_device()_os设置gpu使用设备