How to Configure GPU Support for PyTorch and TensorFlow in AI for Beginners Notebooks

To configure GPU support in the AI for Beginners notebooks, install the NVIDIA driver, CUDA 11.8, and cuDNN 8.x, then use tf.config.experimental.set_memory_growth() for TensorFlow and torch.device("cuda") for PyTorch to move models and tensors onto the GPU.

The Microsoft AI-For-Beginners repository provides a comprehensive curriculum for machine learning that runs significantly faster with GPU acceleration. While the environment.yml already pins GPU-compatible versions of TensorFlow 2.17.0 and PyTorch, you must properly configure device selection and memory management within the notebooks to utilize the hardware. This guide shows the exact code patterns used in the course notebooks and the system prerequisites required to enable GPU support.

System-Level Prerequisites

Before running the notebooks, verify that your host machine meets these requirements. The repository assumes CUDA 11.8 compatibility as specified in environment.yml.

Component	Purpose	Recommended Version	Verification Command
NVIDIA Driver	Low-level GPU access	≥ 450 for CUDA 11.x	nvidia-smi
CUDA Toolkit	Runtime for TensorFlow and PyTorch binaries	11.8	nvcc --version
cuDNN	Deep learning primitives	8.x (matching CUDA)	Check <CUDA_PATH>/include/cudnn.h
Python Packages	Framework binaries with CUDA support	TensorFlow 2.17.0, PyTorch (conda CUDA build)	python -c "import torch; print(torch.version.cuda)"

Note: After installing the NVIDIA software stack, create the ai4beg conda environment using the repository's environment.yml. This automatically pulls the GPU-enabled builds of both frameworks.

Configuring GPU Support in TensorFlow

TensorFlow automatically detects GPUs, but it allocates all available GPU memory by default, which can cause out-of-memory errors when running multiple experiments. The course notebooks use a specific pattern to enable memory growth.

Enabling Memory Growth and Device Detection

The notebooks TransferLearningTF.ipynb and RNNTF.ipynb demonstrate how to list physical devices and configure memory growth before initializing models. Use this snippet at the start of any TensorFlow notebook:

import tensorflow as tf

# List available GPUs

gpus = tf.config.list_physical_devices('GPU')
print("GPUs detected:", gpus)

# Enable memory growth to prevent full allocation

if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
        logical_gpus = tf.config.list_logical_devices('GPU')
        print(f"Enabled memory growth on {len(gpus)} physical GPU(s), "
              f"{len(logical_gpus)} logical GPU(s) available.")
    except RuntimeError as e:
        print(e)  # Memory growth must be set before GPU initialization

After executing this block, subsequent tf.keras operations in lessons such as AutoencodersTF.ipynb automatically execute on the GPU.

Configuring GPU Support in PyTorch

Unlike TensorFlow, PyTorch requires explicit device management. You must manually move models and tensors to the GPU using the .to() method or the device= argument.

Device Selection and Tensor Transfer

The TransferLearningPyTorch.ipynb and IntroPyTorch.ipynb notebooks implement this pattern. Insert this code at the beginning of your PyTorch notebooks:

import torch

# Detect GPU availability

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Move model to GPU

model = torch.nn.Sequential(
    torch.nn.Conv2d(3, 16, kernel_size=3, padding=1),
    torch.nn.ReLU(),
    torch.nn.Flatten(),
    torch.nn.Linear(16 * 32 * 32, 10)
).to(device)  # Model resides on GPU

# Create tensors directly on GPU

sample_input = torch.randn(8, 3, 32, 32, device=device)
output = model(sample_input)  # Computation runs on GPU

Always ensure labels and loss values are also located on the same device as your model weights to avoid device mismatch errors.

GPU Implementation Examples in the Repository

The AI for Beginners course already implements these patterns across multiple lessons. Reference these specific files to see GPU configuration in context:

• lessons/4-ComputerVision/08-TransferLearning/TransferLearningTF.ipynb: Demonstrates TensorFlow GPU detection and memory growth configuration.
• lessons/4-ComputerVision/08-TransferLearning/TransferLearningPyTorch.ipynb: Shows PyTorch device selection and model transfer using .to(device).
• lessons/5-NLP/16-RNN/RNNTF.ipynb: Contains RNN-specific GPU memory management for TensorFlow.
• lessons/3-NeuralNetworks/05-Frameworks/IntroPyTorch.ipynb: Provides a full tutorial on PyTorch GPU usage including tensor device placement.
• lessons/4-ComputerVision/09-Autoencoders/AutoencodersTF.ipynb: Includes GPU availability checks with user warnings.

Troubleshooting Common GPU Configuration Issues

Symptom	Root Cause	Solution
NotFoundError: The source TensorFlow binary is compiled without GPU support	Installed CPU-only TensorFlow build	Use conda install tensorflow==2.17.0 from the environment.yml or ensure PyPI install matches CUDA 11.8
torch.cuda.is_available() returns False	PyTorch CUDA build mismatch with driver	Recreate environment with conda install pytorch torchvision torchaudio cudatoolkit=11.8 -c pytorch
Out-of-memory crashes after several cells	TensorFlow memory growth disabled or PyTorch cache fragmentation	Enable set_memory_growth for TensorFlow; in PyTorch use torch.cuda.empty_cache() and del unused tensors
"Skipping registering GPU devices…" in logs	Container or remote environment without GPU passthrough	Use --gpus all flag in Docker or enable GPU in VS Code Dev Container settings

Quick Setup Checklist for Any Notebook

Add these two cells at the top of any notebook to ensure GPU support is configured correctly:

# TensorFlow GPU setup

import tensorflow as tf
if tf.config.list_physical_devices('GPU'):
    for gpu in tf.config.list_physical_devices('GPU'):
        tf.config.experimental.set_memory_growth(gpu, True)
    print("TensorFlow GPU configured")

# PyTorch GPU setup

import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"PyTorch running on {device}")

Summary

• Install NVIDIA driver ≥ 450, CUDA 11.8, and cuDNN 8.x before creating the conda environment.
• Use tf.config.experimental.set_memory_growth() in TensorFlow notebooks to prevent memory overallocation; the pattern appears in TransferLearningTF.ipynb and RNNTF.ipynb.
• Use torch.device("cuda") and .to(device) in PyTorch to explicitly move models and tensors, as shown in TransferLearningPyTorch.ipynb.
• Verify GPU availability immediately after setup to catch driver or CUDA mismatches before running lengthy training loops.

Frequently Asked Questions

How do I check if my notebook is actually using the GPU?

For TensorFlow, run print(tf.config.list_physical_devices('GPU')) and verify the list is non-empty. For PyTorch, check torch.cuda.is_available() returns True and monitor GPU utilization with nvidia-smi in a separate terminal window while training.

Why does TensorFlow allocate all my GPU memory even for small models?

TensorFlow's default behavior reserves the entire GPU memory address space to prevent fragmentation. According to the source code in TransferLearningTF.ipynb, you must call tf.config.experimental.set_memory_growth(gpu, True) on each physical device before initializing any tensors or models.

Can I use the AI for Beginners notebooks on a CPU-only machine?

Yes. The notebooks gracefully fall back to CPU execution. In PyTorch, torch.device("cuda" if torch.cuda.is_available() else "cpu") automatically selects CPU when no GPU is detected. TensorFlow operations will run on the CPU without modification, though training times will be significantly longer for deep learning lessons.

What should I do if nvidia-smi shows my GPU but PyTorch says CUDA is not available?

This indicates a version mismatch between the PyTorch binary and your NVIDIA driver. The environment.yml in the repository specifies CUDA 11.8. Ensure your driver supports this version and reinstall PyTorch using the conda command provided in the troubleshooting section to ensure you receive the correct CUDA-linked binary.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}