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

> Learn how to configure GPU support for PyTorch and TensorFlow in AI for Beginners notebooks. Install NVIDIA drivers, CUDA, and cuDNN for faster AI model training.

- Repository: [Microsoft/AI-For-Beginners](https://github.com/microsoft/AI-For-Beginners)
- Tags: how-to-guide
- Published: 2026-08-09

---

**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`](https://github.com/microsoft/AI-For-Beginners/blob/main/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`](https://github.com/microsoft/AI-For-Beginners/blob/main/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`](https://github.com/microsoft/AI-For-Beginners/blob/main/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:

```python
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:

```python
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`](https://github.com/microsoft/AI-For-Beginners/blob/main/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:

```python

# 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")

```

```python

# 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`](https://github.com/microsoft/AI-For-Beginners/blob/main/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.