Torchvision fails when trying to use the gpu

Running on Windows 11, installed using the pip command on the website.
CUDA 1.8

I am able to run using the CPU but when I set device=‘0’ it crashes giving me the following error:

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, AutogradMeta, Tracer, AutocastCPU, AutocastCUDA, FuncTorchBatched, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PythonDispatcher].

Running on Windows 11

Verifying pytorch in my console:

PS C:\Users\Dusti> python
Python 3.10.10 (tags/v3.10.10:aad5f6a, Feb  7 2023, 17:20:36) [MSC v.1929 64 bit (AMD64)] on win32
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> torch.rand(3)
tensor([0.2497, 0.1301, 0.3781])
>>> torch.cuda.is_available()
True
>>> torch.cuda.device_count()
1
>>> torch.cuda.current_device()
0
>>> torch.cuda.device(0)
<torch.cuda.device object at 0x0000017B6D0104F0>
>>> torch.cuda.get_device_name(0)
'NVIDIA GeForce RTX 3070 Ti'

My script:

import cv2

from ultralytics import YOLO

def main():
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 780)

    model = YOLO("yolov8l.pt")

    while True:
        ret, frame = cap.read()
        cv2.imshow("yolov8", frame)

        result = model(frame, device='0')

        if(cv2.waitKey(30) == 27):
            break


if __name__ == "__main__":
    main()

Did you install a torchvision version with CUDA support?
If so, what does torchvision.version.cuda return?

I did use this command (generated by selections on the site) to install it all:

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

Oh, maybe the command does not install the cuda version of torchvision?

When I type that command (seems to import torchvision, but the cuda bit fails):

>>> import torchvision
>>> torchvision.version.cuda
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: module 'torchvision.version' has no attribute 'cuda'
>>>

I did a pip list and it looks like it’s not a cuda version

torch              2.0.0+cu118
torchaudio         2.0.1+cu118
torchvision        0.15.1
tqdm               4.65.0
typing_extensions  4.5.0
tzdata             2023.3
ultralytics        8.0.61
urllib3            1.26.15
wheel              0.40.0

any recommendations for the best way to find a cuda version? or commands to uninstall / reinstall a cuda version of the version I have (0.15.1)?

Your install command should work and also installs the proper libraries in my setup:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
Looking in indexes: https://download.pytorch.org/whl/cu118
Collecting torch
  Downloading https://download.pytorch.org/whl/cu118/torch-2.0.0%2Bcu118-cp38-cp38-linux_x86_64.whl (2267.3 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.3/2.3 GB 56.7 MB/s eta 0:00:00
Collecting torchvision
  Downloading https://download.pytorch.org/whl/cu118/torchvision-0.15.1%2Bcu118-cp38-cp38-linux_x86_64.whl (33.9 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 33.9/33.9 MB 57.9 MB/s eta 0:00:00
Collecting torchaudio
  Downloading https://download.pytorch.org/whl/cu118/torchaudio-2.0.1%2Bcu118-cp38-cp38-linux_x86_64.whl (4.4 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 4.4/4.4 MB 54.0 MB/s eta 0:00:00
...
pip list | grep torch
torch              2.0.0+cu118
torchaudio         2.0.1+cu118
torchvision        0.15.1+cu118

Could you uninstall torchvision and reinstall it?

Hmm, it doesn’t seem to install the cuda version when I do a pip install torchvision
I event tried to do a pip install torchvision==0.15.1+cu118

Uninstalled it:

  Successfully uninstalled torchvision-0.15.1
PS C:\Users\Dusti> pip install torchvision==0.15.1+cu118
ERROR: Could not find a version that satisfies the requirement torchvision==0.15.1+cu118 (from versions: 0.1.6, 0.1.7, 0.1.8, 0.1.9, 0.2.0, 0.2.1, 0.2.2, 0.2.2.post2, 0.2.2.post3, 0.12.0, 0.13.0, 0.13.1, 0.14.0, 0.14.1, 0.15.0, 0.15.1)
ERROR: No matching distribution found for torchvision==0.15.1+cu118

I did find the versions on the website
https://download.pytorch.org/whl/torchvision/

Any recommendation?
cp310, cp311, cp38, cp39?

Looks like for you above it downloaded cp38

This tag corresponds to your used Python version and since I’m using Python 3.8 the cp38 wheels were picked.

Perfect!

Thanks for clearing that up. Installed the cp310 version and everything works good.