cuda pytorch

Setting up PyTorch and CUDA

Prerequisites

The following was tested on a ThinkPad P14s with nVidia Quadro T500 GPU under Kubuntu 24.04

$ sudo apt install nvidia-driver-570-server nvidia-utils-570-server nvidia-prime

Confirm successful driver installation and check CUDA version:

$ nvidia-smi

+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 570.133.20             Driver Version: 570.133.20     CUDA Version: 12.8     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA T500                    Off |   00000000:01:00.0 Off |                  N/A |
| N/A   48C    P8            N/A  / 5001W |       5MiB /   4096MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+

[...]

Note that the GPU is switched off in the above summary with prime-select intel. You see the exact driver version (570.133.20) and the CUDA version (12.8).

Configuring a PyTorch conda environment

Specify a minimal conda environment:

name: pytorch
dependencies:
    - python=3.12.*
    - pip

At the time of writing (2025/06), PyTorch does not yet officially support CUDA v12.8:

$ conda search pytorch
[...]
pytorch                        2.7.1 cuda126_mkl_py312_h30b5a27_300  conda-forge
[...]

Install and activate the above environment and install a PyTorch nightly build:

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

After successful installation, you can test your setup as follows:

# run this in an interactive python session
>>> import sys
'3.12.11 | packaged by conda-forge
>>> import torch
>>> torch.__version__
'2.7.1+cu128'
>>> torch.cuda.is_available()
True
>>> torch.cuda.device_count()
1
>>> torch.cuda.current_device()
0
>>> torch.cuda.device(0)
<torch.cuda.device object at 0x7982c7f4fd70>
>>> torch.cuda.get_device_name(0)
'NVIDIA T500'
>>>