我正在尝试运行一个简单的 pytorch 示例代码。使用 CPU 效果很好。但是在使用 GPU 时,我收到此错误消息:
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/conv.py", line 263, in forward
return self._conv_forward(input, self.weight, self.bias)
File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/conv.py", line 260, in _conv_forward
self.padding, self.dilation, self.groups)
RuntimeError: cuDNN error: CUDNN_STATUS_NOT_INITIALIZED
我试图运行的代码如下:
import torch
from torch import nn
m = nn.Conv1d(16, 33, 3, stride=2)
m=m.to('cuda')
input = torch.randn(20, 16, 50)
input=input.to('cuda')
output = m(input)
我在带有 CUDA 10.2 版的 NVIDIA docker 中运行这段代码,我的 GPU 是 RTX 2070
原文由 Eduardo H 发布,翻译遵循 CC BY-SA 4.0 许可协议
这里 有一些关于这个的讨论。我有同样的问题,但使用 cuda 11.1 为我解决了它。
这是确切的 pip 命令