我尝试在我的笔记本电脑上使用PyCharm在Windows 10上使用GPU训练模型，所以我在网上遵循了一些关于如何正确设置我的NVidia驱动程序，cuda，cuDNN以及安装库版本的指南。显然安装过程是成功的，并且能够发现GPU并获得有关它的信息。然而，如果我尝试训练模型，我会得到错误：

NotImplementedError: Could not run 'torchvision::nms with arguments from the 'CUDA' backend.'

字符串
我已经设置了 Torch ， Torch 视觉和 Torch 音频的指示，对他们的website

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

型
我已经安装了最新的NVidia驱动程序，CUDA版本12.1和cuDNN，并确保环境变量中的路径正确。我重新启动了我的电脑，仍然什么都没有。
如果我尝试检查我的GPU信息，它会按预期工作。

import torch

print(f'CUDA version: {torch.version.cuda}')
cuda_id = torch.cuda.current_device()
print(f"ID of current CUDA device:{torch.cuda.current_device()}")
print(f"Name of current CUDA device:{torch.cuda.get_device_name(cuda_id)}")
print(f'Torch version: {torch.__version__}')

型
输出量：

CUDA version: 12.1
ID of current CUDA device:0
Name of current CUDA device:NVIDIA GeForce GTX 1060 with Max-Q Design
Torch version: 2.1.1+cu121

型
所以Torch版本是正确的，并且与我的CUDA版本相对应。
我尝试训练的代码：

from ultralytics import YOLO
from datetime import datetime, timedelta

model = YOLO('yolov8n.pt')

startTime = datetime.now()
results = model.train(data='data.yaml', epochs=15,
                      imgsz=[480, 352])
endTime = datetime.now()
delta = endTime - startTime
print('\n\n\n')
print(f'Model training took {delta}')
print(f'Model training took {delta.seconds}')

型
这还是安装不正确的问题吗？可能是pyCharm造成的吗？我不知所措。