How to configure PyTorch with GPU support

What this does

Installs the PyTorch package with GPU acceleration and verifies that the framework can detect and allocate work on the available accelerator. This is the foundation for running transformer models and inference servers on local hardware.

Steps

1. Create and activate a virtual environment.

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

   Expected output: Prompt reflects the activated environment.

2. Install PyTorch with the GPU-enabled index.

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

   Expected output: Wheels downloaded and installed.

3. Validate the installation.

   python3 -c "import torch; print('CUDA available:', torch.cuda.is_available()); print('Device:', torch.cuda.get_device_name(0))"
   

   Expected output: True on first line, GPU model string on second line.

4. Execute a minimal matrix multiplication on the GPU.

   python3 -c "
   import torch
   a = torch.randn(4096, 4096, device='cuda')
   b = torch.randn(4096, 4096, device='cuda')
   c = a @ b
   print('GPU compute succeeded, result shape:', c.shape)
   "
   

   Expected output: GPU compute succeeded, result shape: torch.Size([4096, 4096])

Verification

python3 -c "import torch; print('CUDA:', torch.cuda.is_available(), 'Count:', torch.cuda.device_count(), 'Name:', torch.cuda.get_device_name(0))"
# Expected: CUDA: True Count: 1 Name: <GPU model>

Common failures

• torch.cuda.is_available() returns False — The PyTorch build does not match host CUDA version. Reinstall from the correct index URL.
• CUDA error: no kernel image is available for execution — NVIDIA driver is older than the CUDA toolkit used to build PyTorch. Upgrade the driver.
• No module named 'torch' — Run install inside the activated virtual environment.
• cuDNN not detected — Ensure libcudnn8 is installed and visible to Python.
• Out-of-memory error — Reduce tensor size or close other GPU processes. Run nvidia-smi to inspect memory.

Related guides

• How to verify CUDA setup for AI workloads
• How to install NVIDIA Container Toolkit for Docker
• Course Local AI Fundamentals