RuntimeError: CUDA error: unknown error in torch 1.1.0

1

I have some problems when I try to use GPU.
Here is my code from example:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
densenet121 = models.densenet121(pretrained=True)
densenet121.to(device)

When I run the code above, I get the error:

Traceback (most recent call last):
  File "d:\All Project\VScode\Python\learn_framework\CNN\test.py", line 63, in <module>
    densenet121.to(device)
  File "D:\Python36\lib\site-packages\torch\nn\modules\module.py", line 386, in to
    return self._apply(convert)
  File "D:\Python36\lib\site-packages\torch\nn\modules\module.py", line 193, in _apply
    module._apply(fn)
  File "D:\Python36\lib\site-packages\torch\nn\modules\module.py", line 193, in _apply
    module._apply(fn)
  File "D:\Python36\lib\site-packages\torch\nn\modules\module.py", line 199, in _apply
    param.data = fn(param.data)
  File "D:\Python36\lib\site-packages\torch\nn\modules\module.py", line 384, in convert
    return t.to(device, dtype if t.is_floating_point() else None, non_blocking)
  File "D:\Python36\lib\site-packages\torch\cuda\__init__.py", line 163, in _lazy_init
    torch._C._cuda_init()
RuntimeError: CUDA error: unknown error

But if I run code like this:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.randn(1).to(device)
densenet121 = models.densenet121(pretrained=True)
densenet121.to(device)
print(next(densenet121.parameters()).is_cuda)

It will work and print a ‘True’.
I don’t know what the problem is and don’t want to add ‘torch.randn(1).to(device)’ in front of the code every time.
My machine:
Windows 10 Enterprise, 64-bit
nvidia driver: 418.96
GPU: GeForce 940MX
CUDA: release 10.0, V10.0.130
Pytorch version: 1.1.0

Please help:pleading_face:

2