This guide walks through the setup of PyTorch across three common environments: Windows (with CUDA), NVIDIA Jetson Nano (JetPack 4.6), and Ubuntu Linux. It covers driver installation, CUDA Toolkit, cuDNN, and final verification.

1. Windows Installation

---

1.1 Verify or Install NVIDIA Driver

The NVIDIA driver acts as the communication bridge between the OS and the GPU hardware. CUDA operations require a compatible driver. Open a command prompt (Win+R, type cmd) and run:

nvidia-smi

If the command is not recognized, download and install the latest driver from the NVIDIA Driver Download page. Note the displayed CUDA version (e.g., 12.5). This version indicates the maximum CUDA level your driver supports.

1.2 Install CUDA Toolkit

CUDA Toolkit provides the compiler, runtime, and libraries for GPU computing. Choose a version equal to or lower than the driver's CUDA version. For a driver showing CUDA 12.5, download CUDA Toolkit 12.1.0 (or the matching version).

1. Visit the CUDA Toolkit download page.
2. Select Windows, x86_64, 11 (for Windows 11/10), and exe (local).
3. Run the installer. Choose Custom (Advanced). The first path is a temporary extraction folder; the second is the actual installation diretcory (e.g., D:\NVIDIA\CUDA\v12.1). Note this location for later steps.

After installation, verify with:

nvcc -V

1.3 Install cuDNN

cuDNN is a GPU‑accelerated library for deep neural networks. It provides highly optimized implementations of convolution, pooling, normalization, etc.

1. Download cuDNN from the NVIDIA cuDNN page (requires free registration). Choose the version matching your CUDA (e.g., for CUDA 12.x).
2. Extract the archive. Inside you will find bin, include, lib folders.
3. Copy these three folders into the CUDA installation directory (e.g., D:\NVIDIA\CUDA\v12.1), accepting the merge.

1.4 Create a Conda Environment and Install PyTorch

Open Anaconda Prompt and create a new environment (optional):

conda create -n pytorch_env python=3.9
conda activate pytorch_env

Visit the PyTorch website, select your configuration (Stable, Windows, Conda, CUDA 12.1), and copy the generated command. For example:

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

Run the command inside the activated environment. Wait for the download and installation to complete.

1.5 Verify Installation

python -c "import torch; print('CUDA available:', torch.cuda.is_available()); print('PyTorch version:', torch.__version__)"

Expect True for CUDA availability and the version string.

2. Jetson Nano (JetPack 4.6, Ubuntu 18.04)

---

Environment: Ubuntu 18.04, JetPack 4.6 (Python 3.6, CUDA 10.2). The Nano runs on ARM64, so pre‑built wheels may not work. We install PyTorch 1.8.0 and torchvision 0.9.0 from source.

2.1 Update System and Install Dependencies

sudo apt-get update
sudo apt-get install -y python-pip libopenblas-base libopenmpi-dev libomp-dev \
  build-essential make cmake cmake-curses-gui git g++ pkg-config curl \
  libatlas-base-dev gfortran libcanberra-gtk-module libcanberra-gtk3-module \
  libhdf5-serial-dev hdf5-tools nano locate screen \
  libfreetype6-dev protobuf-compiler libprotobuf-dev openssl \
  libssl-dev libcurl4-openssl-dev cython3 \
  libjpeg-dev zlib1g-dev libpython3-dev libavcodec-dev libavformat-dev libswscale-dev \
  libopenmpi2 libopenblas-dev

Install additional Python packages:

pip install Cython
pip install pillow==8.4.0   # avoid too‑new pillow

2.2 Set CUDA Environment Variables

echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc
nvcc --version

2.3 Install PyTorch 1.8.0

Download the pre‑built wheel for JetPack 4.6 from NVIDIA’s forum (e.g., torch-1.8.0-cp36-cp36m-linux_aarch64.whl). Install with:

pip install ~/Downloads/torch-1.8.0.whl

2.4 Build torchvision 0.9.0 from Source

git clone --branch v0.9.0 https://github.com/pytorch/vision torchvision
cd torchvision
export BUILD_VERSION=0.9.0
# If you get errors about ffmpeg, edit setup.py and change 'if has_ffmpeg' to 'if False'
sudo python3 setup.py install --user

2.5 Verify

python -c "import torch; print(torch.cuda.is_available()); print(torch.__version__); import torchvision; print(torchvision.__version__)"

Should print True, 1.8.0, 0.9.0.

3. Ubuntu Linux (Desktop / Server)

---

3.1 Check or Install NVIDIA Driver

nvidia-smi

If missing, install via the Software & Updates GUI (Additional Drivers tab) or using the command line:

sudo apt-get install nvidia-driver-535   # or the recommended version
sudo reboot

3.2 Install CUDA Toolkit

Download the runfile (or deb) from the NVIDIA CUDA Toolkit archive. For example, for CUDA 12.4:

wget https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run
sudo sh cuda_12.4.0_550.54.14_linux.run

During installation, uncheck the driver option if you already have a compatible one. Add the following lines to ~/.bashrc:

export PATH=/usr/local/cuda-12.4/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-12.4/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

Reload and verify:

source ~/.bashrc
nvcc -V

3.3 Install cuDNN

Download the cuDNN archive for CUDA 12.x (e.g., cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar.xz). Extract and copy:

xz -d cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar.xz
tar -xf cudnn-linux-x86_64-8.9.7.29_cuda12-archive.tar
cd cudnn-linux-x86_64-8.9.7.29_cuda12-archive
sudo cp include/cudnn*.h /usr/local/cuda-12.4/include
sudo cp lib/libcudnn* /usr/local/cuda-12.4/lib64
sudo chmod a+r /usr/local/cuda-12.4/include/cudnn*.h /usr/local/cuda-12.4/lib64/libcudnn*

3.4 Install PyTorch via Pip or Conda

Create a conda environment (optional) and install PyTorch. Example command for pip with CUDA 12.4:

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

3.5 Verify

python3 -c "import torch; print('CUDA:', torch.cuda.is_available()); print('Version:', torch.__version__)"

3.6 Uninstallation (if needed)

pip uninstall torch torchvision torchaudio
pip show torch    # confirm removal

Common Issues on Jetson Nano

• fatal error: cuda.h – ensure CUDA paths are set and cuda.h is in /usr/local/cuda/include.
• avcodec/avcodec.h missing – in torchvision’s setup.py, change if has_ffmpeg to if False to skip FFmpeg detection.
• Python version mismatch – JetPack 4.6 ships Python 3.6; ensure you use the correct wheel and base Python.