PyTorch Installation

By the end of this lesson, you will be able to install PyTorch with GPU support and verify that your system correctly detects CUDA.

What it is

PyTorch is an open-source machine learning library. While it can run on a CPU, deep learning models train significantly faster on NVIDIA GPUs using CUDA (Compute Unified Device Architecture). Installing PyTorch involves selecting the correct build variant: CPU-only or GPU-enabled. The GPU version requires specific drivers and a matching CUDA toolkit version installed on your operating system. Related terms include cu118, cu121 (CUDA versions), and torch.cuda.is_available().

Why it matters

• Performance: GPU training can be 10x-100x faster than CPU for large matrices.
• Compatibility: Mismatched CUDA versions cause silent failures or runtime errors.
• Resource Management: Correct installation ensures memory is allocated to the GPU, not RAM.
• Reproducibility: Standardized installs help teams share code without environment conflicts.

Syntax or steps

The standard method uses pip with a custom index URL. You must match the CUDA version supported by your NVIDIA driver.

1. Check your NVIDIA driver version using nvidia-smi.
2. Determine the highest CUDA version your driver supports.
3. Select the corresponding PyTorch wheel from the official website.
4. Run the pip command with the --index-url flag.

Example

# Step 1: Install PyTorch with CUDA 12.1 support
# Note: Replace 'cu121' with your required version (e.g., cu118)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# Step 2: Verify installation in Python
import torch

print(f"PyTorch Version: {torch.__version__}")
print(f"CUDA Available: {torch.cuda.is_available()}")

if torch.cuda.is_available():
    print(f"Device Count: {torch.cuda.device_count()}")
    print(f"Current Device: {torch.cuda.current_device()}")
    print(f"Device Name: {torch.cuda.get_device_name(0)}")
else:
    print("Running on CPU only.")

Explanation: The first block installs the libraries. The second block imports torch and checks is_available(). If true, it prints hardware details. If false, it confirms CPU fallback.

Common mistakes

• Driver/CUDA Mismatch: Installing cu121 when your driver only supports up to CUDA 11.8. Fix: Check nvidia-smi header for "CUDA Version".
• Missing Index URL: Running pip install torch defaults to CPU. Fix: Always specify the --index-url for GPU builds.
• Conda vs Pip Conflicts: Mixing Conda and Pip environments can break dependencies. Fix: Use one package manager consistently within a virtual environment.
• Ignoring WSL2: On Windows, native Linux tools often work better than PowerShell for complex setups. Fix: Consider using Windows Subsystem for Linux if issues persist.

When to use it

Compare installation methods based on your needs.

Method	Best For	Complexity
Pip (Official)	Standard development, quick setup	Low
Conda	Data science workflows, managing non-Python deps	Medium
Docker	Production deployment, strict reproducibility	High

Use Pip for most local development tasks. Use Docker if you need to guarantee the exact same environment across different machines.

Practice

Guided Exercise: Run nvidia-smi in your terminal. Note the "CUDA Version" displayed at the top right. Go to pytorch.org and find the install command that matches or is lower than that version. Execute it.

Challenge: Write a script that attempts to move a tensor to the GPU. If it fails, catch the exception and print "GPU not detected."
Hint: Use try...except RuntimeError around tensor.to('cuda').

Quick check

Question: Why does torch.cuda.is_available() return False even after installing PyTorch? Answer: Likely because you installed the CPU-only version (missing --index-url) or your NVIDIA drivers are outdated/incompatible with the installed CUDA toolkit.

Summary

Installing PyTorch with GPU support requires matching your NVIDIA driver's CUDA capability with the correct PyTorch wheel via pip. Always verify success using torch.cuda.is_available() before starting model training to ensure computational resources are properly utilized.