torch.save() gives : RuntimeError: CUDA error: no CUDA-capable device is detected

Question

我在 GPU 上训练神经网络模型，但是当我使用

torch.save()

保存检查点时出现上述错误。我的问题是，即使我有一个 CUDA 设备，为什么我会收到上述错误？我的模型运行在 GPU 上没问题，请查看下面的输出：nvidia-smi 命令。

$ nvidia-smi
Sat Aug 15 09:51:58 2020       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 440.100      Driver Version: 440.100      CUDA Version: 10.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce RTX 2060    Off  | 00000000:01:00.0 Off |                  N/A |
| N/A   55C    P3    33W /  N/A |   4774MiB /  5934MiB |     97%      Default |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0      7080      C   python3                                     4763MiB |
+-----------------------------------------------------------------------------+

$ python --version
Python 3.8.2

$ python -c "import torch; print(torch.__version__)"
1.5.1

$ python -c "import torchvision as torch; print(torch.__version__)"
0.6.1

我什至尝试了以下方法：

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

torch.save({
                'epoch': epoch + 1,
                'metrics': metrics,
                'model': model.state_dict(),
                'optimizer' : optimizer.state_dict(),
            }, name)

但没有任何效果。我是深度学习新手，仍在学习 PyTorch。请原谅我的无知。

Answer 1

好像是其他问题。测试这个玩具示例，看看问题是否出在 torch.save() 或任何其他命令中。

import torch
import torchvision

model = torchvision.models.resnet18()
model = model.cuda()

torch.save(model.state_dict(), 'net')

torch.save() gives : RuntimeError: CUDA error: no CUDA-capable device is detected

torch.save() gives : RuntimeError: CUDA error: no CUDA-capable device is detected

deep-learning

pytorch

torchvision