Search code examples
pythonpytorchyolo

Getting "NotImplementedError: Could not run 'torchvision::nms' with arguments from CUDA backend" despite having all necessary libraries and imports


The full error:

NotImplementedError: Could not run 'torchvision::nms' with arguments from the 'CUDA' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'torchvision::nms' is only available for these backends: [CPU, QuantizedCPU, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, AutogradMPS, AutogradXPU, AutogradHPU, AutogradLazy, Tracer, AutocastCPU, AutocastCUDA, FuncTorchBatched, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PythonDispatcher].

I get this when attempting to train a YOLOv8 model on a Windows 11 machine, everything works for the first epoch then this occurs.


I also get this error immediately after the first epoch ends but I don't think it is relevant.

Error executing job with overrides: ['task=detect', 'mode=train', 'model=yolov8n.pt', 'data=custom.yaml', 'epochs=300', 'imgsz=160', 'workers=8', 'batch=4']

I was trying to train a YOLOv8 image detection model utilizing CUDA GPU.


Solution

  • This error occurs when you have Torch and Torchaudio for CUDA but not Torchvision for CUDA installed.

    Uninstall Torch and Torchvision, I used pip:

    pip uninstall torch torchvision
    

    Then go here to install the proper versions of both using the nice interface. I got the following command:

    pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116