Pytorch CUDA报错，RuntimeError: cuDNN error: CUDNN_STATUS_INTERNAL_ERROR

（多卡环境）在使用torchsummary()进行可视化的时候，代码报错：

RuntimeError: cuDNN error: CUDNN_STATUS_INTERNAL_ERROR

在代码中设置指定GPU，代码仍旧运行不成功

device = torch.device(“cuda:3” if torch.cuda.is_available() else “cpu”)
model = model.to(device)

import torch
from torchsummary import summary
from torchvision.models import vgg11

torch.cuda.set_device(2)

model = vgg11(pretrained=False)
if torch.cuda.is_available():
    # device = torch.device("cuda:3")
    # model = model.to(device)
    model.cuda()
summary(model, (3, 224, 224))

添加了代码torch.cuda.set_device(2)设置当前设备。然后运行成功，其结果如下图所示：

实际上，官方文档提示不鼓励使用此函数来设置。在大多数情况下，最好使用CUDA_VISIBLE_DEVICES环境变量。

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

import torch
from torchsummary import summary
from torchvision.models import vgg11

model = vgg11(pretrained=False)
if torch.cuda.is_available():
    model.cuda()
summary(model, (3, 224, 224))

注意

1. CUDA_VISIBLE_DEVICES设置要在模型加载到GPU上之前
2. 使用os.environ['CUDA_VISIBLE_DEVICES']对可以使用的显卡进行限定之后, 显卡的实际编号和程序看到的编号应该是不一样的, 例如上面我们设定的是os.environ['CUDA_VISIBLE_DEVICES']="0,2", 但是程序看到的显卡编号应该被改成了'0,1' 也就是程序所使用的显卡编号实际上是经过了一次映射之后才会映射到真正的显卡编号上面的。