报错复现:
Traceback (most recent call last):
File "../train.py", line 388, in <module>
main()
File "../train.py", line 384, in main
train(args)
File "../train.py", line 378, in train
trainer.train()
File "../train.py", line 322, in train
self._init_model()
File "../train.py", line 308, in _init_model
self._set_device()
File "../train.py", line 168, in _set_device
self._model.to(self._device)
File "D:\Apps\Anaconda\envs\tse_2021_bases\lib\site-packages\torch\nn\modules\module.py", line 425, in to
return self._apply(convert)
File "D:\Apps\Anaconda\envs\tse_2021_bases\lib\site-packages\torch\nn\modules\module.py", line 201, in _apply
module._apply(fn)
File "D:\Apps\Anaconda\envs\tse_2021_bases\lib\site-packages\torch\nn\modules\module.py", line 223, in _apply
param_applied = fn(param)
File "D:\Apps\Anaconda\envs\tse_2021_bases\lib\site-packages\torch\nn\modules\module.py", line 423, in convert
return t.to(device, dtype if t.is_floating_point() else None, non_blocking)
File "D:\Apps\Anaconda\envs\tse_2021_bases\lib\site-packages\torch\cuda\__init__.py", line 196, in _lazy_init
_check_driver()
File "D:\Apps\Anaconda\envs\tse_2021_bases\lib\site-packages\torch\cuda\__init__.py", line 94, in _check_driver
raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled
可见在将model送到GPU的过程中报错,CUDA不可用。
解决:
import torch
print(torch.__version__)
print(torch.cuda.is_available())
通过以上代码查看已安装pytorch版本是否可以调用GPU
返回为False的话,去pytorch官网依照命令重新安装相应版本的pytorch
pytorch与cuda版本问题参考链接。