Install PyTorch with GPU Support

PyTorch's GPU support can dramatically accelerate training and inference, but getting the right build—CUDA for NVIDIA hardware, ROCm for AMD—against your actual driver stack is where installs go wrong. This guide walks through both paths cleanly, using a virtual environment so system Python stays untouched.

Before You Start: Know Your Hardware

The build of PyTorch you install must match your GPU vendor and the compute runtime installed on the host. Mixing these is the single most common source of silent CPU fallback.

• NVIDIA GPU: You need the proprietary NVIDIA driver and the CUDA toolkit (or at minimum the CUDA runtime). The PyTorch wheel bundles its own libcudart, so you do not need the full CUDA toolkit installed—but your driver must support the CUDA version PyTorch was compiled against.
• AMD GPU: You need ROCm installed from AMD's repositories. ROCm support is Linux-only and limited to specific GPU generations (RDNA2/CDNA and newer are best supported).
• CPU only: If you have no supported GPU or just want to prototype, a CPU-only wheel is the safe fallback.

Check Your NVIDIA Driver and CUDA Version

nvidia-smi

The top-right corner of the output shows the maximum CUDA version your driver supports (e.g., CUDA Version: 12.4). Choose a PyTorch build at or below that version.

Check Your AMD ROCm Version

rocminfo | grep -i 'ROCm'

Or check the installed package version directly:

# Debian/Ubuntu
apt-cache show rocm-dev | grep Version

# Fedora/RHEL family
dnf info rocm-dev

Step 1: Create and Activate a Virtual Environment

Never install PyTorch into the system Python. A virtual environment gives you clean reproducibility and easy teardown.

python3 -m venv ~/pytorch-env
source ~/pytorch-env/bin/activate

Confirm you are inside the venv—your prompt should be prefixed with (pytorch-env), and which python should point into your home directory.

which python
# Example output: /home/youruser/pytorch-env/bin/python

Upgrade pip before installing anything; older pip versions can pick the wrong wheel for your platform.

pip install --upgrade pip

Step 2: Install PyTorch — NVIDIA CUDA Path

The PyTorch "Get Started" page generates the exact install command for your configuration. The commands below reflect current stable releases; always verify against the official page for the latest version.

CUDA 12.1 build (covers drivers supporting CUDA 12.x)

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

CUDA 11.8 build (older drivers)

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

The --index-url flag points pip at PyTorch's own wheel index rather than PyPI. Without it you may silently receive the CPU-only build. torchvision and torchaudio are optional but recommended; their versions must match the PyTorch release.

Step 3: Install PyTorch — AMD ROCm Path

ROCm support requires ROCm 5.6 or later on the host. PyTorch ships ROCm wheels for Linux x86-64 only.

Install ROCm on Debian/Ubuntu (if not already present)

wget https://repo.radeon.com/amdgpu-install/latest/ubuntu/jammy/amdgpu-install_6.1.60100-1_all.deb
sudo apt install ./amdgpu-install_6.1.60100-1_all.deb
sudo amdgpu-install --usecase=rocm
sudo usermod -aG render,video $USER

Log out and back in after adding yourself to the render and video groups. Skipping this step causes permission errors when PyTorch tries to open the GPU device.

Install ROCm on Fedora/RHEL family

sudo dnf install https://repo.radeon.com/amdgpu-install/latest/rhel/9.3/amdgpu-install-6.1.60100-1.el9.noarch.rpm
sudo amdgpu-install --usecase=rocm
sudo usermod -aG render,video $USER

Install the ROCm PyTorch wheel

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

Match the rocm version suffix to your installed ROCm version. PyTorch's wheel index currently offers rocm5.6 and rocm5.7; check the index URL directly if you need another: https://download.pytorch.org/whl/.

Step 4: CPU-Only Install (Fallback)

If you have no supported GPU or are provisioning a CI machine:

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

Step 5: Verify GPU Access

This is the critical check. Run it immediately after installing.

python - <<'EOF'
import torch
print("PyTorch version:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
print("CUDA version built against:", torch.version.cuda)
if torch.cuda.is_available():
    print("GPU:", torch.cuda.get_device_name(0))
    t = torch.tensor([1.0, 2.0]).cuda()
    print("Tensor on GPU:", t)
EOF

Expected output with a working NVIDIA setup looks similar to:

# PyTorch version: 2.3.0+cu121
# CUDA available: True
# CUDA version built against: 12.1
# GPU: NVIDIA GeForce RTX 3080
# Tensor on GPU: tensor([1., 2.], device='cuda:0')

For ROCm / AMD GPUs

ROCm exposes itself through the same CUDA API surface in PyTorch, so torch.cuda.is_available() returns True on a working ROCm setup as well. You can also check:

python -c "import torch; print(torch.version.hip)"

A non-None result confirms the wheel was built with ROCm (HIP) support.

Step 6: Freeze Your Environment

Once everything works, pin your dependencies so the environment is reproducible.

pip freeze > requirements.txt

To recreate the environment on another machine, activate a fresh venv then:

pip install -r requirements.txt --index-url https://download.pytorch.org/whl/cu121

Include the --index-url flag here too, or pip will not find the GPU wheels.

Troubleshooting

torch.cuda.is_available() returns False

• Confirm nvidia-smi works outside the venv. If it errors, the driver is not loaded—check dmesg | grep -i nvidia and reinstall the driver.
• Verify the wheel's CUDA version is not higher than your driver supports. A CUDA 12.4 wheel will not work on a driver capped at CUDA 11.8.
• Make sure you did not accidentally install the CPU wheel. Run pip show torch and check the version string; it should end in +cu121 (or similar), not just the version number alone.

ROCm: No GPU found or HSA_STATUS_ERROR_OPEN_FAILED

• Your user is not in the render and video groups, or you have not logged out since adding them. Verify with groups.
• Run rocminfo as your normal user. If it fails with a permission error, the group membership issue is confirmed.

Slow first run / kernel compilation

On first use, CUDA JIT-compiles kernels and caches them under ~/.cache/torch. This is normal and only happens once per configuration. Set TORCH_COMPILE_DEBUG=1 if you need to diagnose compilation failures.

Out of memory errors

PyTorch does not release GPU memory to the OS immediately. Call torch.cuda.empty_cache() between large allocations in a script, or reduce your batch size. Check current allocation with:

python -c "import torch; print(torch.cuda.memory_summary())"