torch.save给出：RuntimeError：CUDA错误：未检测到具有CUDA功能的设备

如何解决torch.save给出：RuntimeError：CUDA错误：未检测到具有CUDA功能的设备

我当时正在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。请原谅我的无知。

解决方法

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

import torch
import torchvision

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

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