CUDA is NVIDIA’s parallel computing platform and programming interface for running general-purpose computations on NVIDIA GPUs. In PyTorch workflows, it accelerates tensor operations, model training, and inference by moving computation from CPU to GPU memory. section 2.2.1 section 3.2.4
Run a Transformers pipeline on CUDA
from transformers import pipeline classifier = pipeline( task="text-classification", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english", device="cuda" ) classifier("The service was excellent and I would gladly return.") The device parameter selects CUDA instead of the default CPU execution. The pipeline returns a classification label and confidence score. section 3.2.4
CUDA support is useful when you need to:
- Accelerate training and inference for deep learning models. section 2.2.1
- Use PyTorch to transfer tensor operations between CPU and GPU memory. section 2.2.1
- Distribute a model across multiple GPUs when it cannot fit on one device. section 2.2.1
- Select CUDA automatically when loading a pretrained model alongside other hardware backends such as Apple MPS. section 6.1.4
$ pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 -U
Install CUDA-compatible PyTorch packages from the cu121 wheel index before testing GPU availability. section 2.2.1
Listing 2.1 Using the torch package to find details on your GPU
import torch
use_cuda = torch.cuda.is_available()
if use_cuda:
print('__CUDNN VERSION:', torch.backends.cudnn.version())
print('__Number CUDA Devices:', torch.cuda.device_count())
print('__CUDA Device Name:', torch.cuda.get_device_name(0))
print('__CUDA Device Total Memory [GB]:',
torch.cuda.get_device_properties(0).total_memory/1e9)
Inspect CUDA availability, cuDNN version, device count, GPU name, and total GPU memory with torch. section 2.2.1
