CUDA Not Detected in PyTorch - Unable to Use GPU for YOLO Training

[AiMliterate]

Hello,
JETSON ORIN NX 16GB
I’m encountering an issue where my system is not detecting CUDA, even though I have installed CUDA 12.6 and PyTorch 2.5.1. When I try to run a YOLOv8 training command, it throws the following error:

Command:

bash

yolo train data=data.yaml model=yolov8s.pt epochs=80 imgsz=640 batch=16 device=0

Error:
ValueError: Invalid CUDA ‘device=0’ requested. Use ‘device=cpu’ or pass valid CUDA device(s) if available, i.e. ‘device=0’ or ‘device=0,1,2,3’ for Multi-GPU.

torch.cuda.is_available(): False
torch.cuda.device_count(): 0
os.environ[‘CUDA_VISIBLE_DEVICES’]: None
See Start Locally | PyTorch for up-to-date torch install instructions if no CUDA devices are seen by torch.

I also ran a small Python script to check CUDA availability in PyTorch, and here’s the result:

Code:
check.py

import torch

Check if CUDA is available

cuda_available = torch.cuda.is_available()

print(f"CUDA Available: {cuda_available}")

If CUDA is available, print the number of GPUs

if cuda_available:

print(f"Number of GPUs available: {torch.cuda.device_count()}")

Print the name of the GPU

print(f"GPU Name: {torch.cuda.get_device_name(0)}")

else:

print(“CUDA is not available on this system.”)

Output:

python check.py

CUDA Available: False

CUDA is not available on this system.

System Info:
JETSON ORIN NX 16GB
jetpack6.1
apt show nvidia-jetpack
Package: nvidia-jetpack
Version: 6.1+b123
Priority: standard
Section: metapackages
Source: nvidia-jetpack (6.1)
Maintainer: NVIDIA Corporation
Installed-Size: 199 kB
Depends: nvidia-jetpack-runtime (= 6.1+b123), nvidia-jetpack-dev (= 6.1+b123)
Homepage: Jetson - Embedded AI Computing Platform | NVIDIA Developer
Download-Size: 29.3 kB
APT-Sources: https://repo.download.nvidia.com/jetson/common r36.4/main arm64 Packages
Description: NVIDIA Jetpack Meta Package

flashed using nvidia SDK manager
CUDA Version:
bash

nvcc --version

Output:
nvcc: NVIDIA (R) Cuda compiler driver

Copyright (c) 2005-2024 NVIDIA Corporation

Built on Wed_Aug_14_10:14:07_PDT_2024

Cuda compilation tools, release 12.6, V12.6.68

Build cuda_12.6.r12.6/compiler.34714021_0

PyTorch Version:
bash

python -c “import torch; print(torch.version)”

Output:

2.5.1

---

Issue:

Even though CUDA 12.6 is installed and recognized by the system, PyTorch is not detecting the CUDA devices. I’m unsure why my GPU is not being recognized by PyTorch, and I’d appreciate any guidance on how to resolve this issue.

Steps I’ve Taken:

1. Verified that the system recognizes the CUDA toolkit with nvcc --version.
2. Checked the PyTorch installation with python -c “import torch; print(torch.version)”.
3. Ran a simple script to verify CUDA availability, but it shows False for torch.cuda.is_available().

Can anyone help me figure out why CUDA isn’t available in PyTorch despite being installed and configured?

[DavidDDD]

Hi

Please download prebuilt Pytorch wheel from this website.

Thanks

[AiMliterate]

Thank you @DavidDDD !

[AiMliterate]

Hi @DavidDDD , i am facing error here can you help me out?

ERROR:

**(modeltraining1) user-name@ubuntu:~/Desktop/vscode/Trial/yolov8_modelTraining$ python t2.py **
No GPU detected. Using CPU for training.

i am training a dataset to get a pytorch file(.pt) using yolo.
i want to run it on GPU… this is my first time using jetson orin nx. so i dont know if it is possible to access gpu and run Epochs in less time.

Thanks in advance!

[DavidDDD]

Hi,

Could you try below commands and share the results for us?

from ultralytics import YOLO
import torch

model = YOLO('yolov5s.pt') # change to your .pt file
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
print(next(model.parameters()).device)
#cuda:0

Thanks

[AiMliterate]

I tried executing the provided code,

saved the .py file in Desktop and got the following output:
**python Tcuda.py **
PRO TIP 💡 Replace ‘model=yolov5s.pt’ with new ‘model=yolov5su.pt’.
YOLOv5 ‘u’ models are trained with “https:github.com\ltralytics/ultralytics” and feature improved performance vs standard YOLOv5 models trained with "https:github.com\ultralytics/yolov5.

Downloading https:/github.com/ultralytics/assets/releases/download/v8.3.0/yolov5su.pt to ‘yolov5su.pt’…
100%|██████████████████████████████████████| 17.7M/17.7M [00:03<00:00, 6.13MB/s]
cuda:0

Also, Executed in vscode:
Tcuda.py
PRO TIP 💡 Replace ‘model=yolov5s.pt’ with new ‘model=yolov5su.pt’.
YOLOv5 ‘u’ models are trained with https:/github.com/ultralytics/ultralytics and feature improved performance vs standard YOLOv5 models trained with https:/github.com/ultralytics/yolov5.

cpu

note:
ignore the url as i had to change it, because unable to mention 4 Links(because i am a new user) in one qurey.

Thanks

[DavidDDD]

Hi,

@AiMliterate

cuda:0

Based on the provided log, both Pytorch and YOLO appear to function correctly.
The issue might be related to VS Code.

Thanks

[AiMliterate]

Hi, @DavidDDD Thank you for clarifying
I now tried to running the script to train the model without using vscode,

When running scripts in VS Code, the environment was not correctly set up to use CUDA. If torch.cuda.is_available() returns False, it means the GPU was not being utilized.

Now it is using gpu.

Thank you.