TensorFlow Setup: Google Colab, Local Install and GPU Check

The Story — Renting a Kitchen vs Building One

In this lesson you will set up TensorFlow in three ways. Then you will check if TensorFlow can see a GPU. This course uses TensorFlow 2.21. It supports Python 3.10 to 3.13.

Pick Your Setup

Option 1 — Google Colab

import tensorflow as tf

print("TensorFlow:", tf.__version__)
print("GPUs:", tf.config.list_physical_devices('GPU'))

You can also see the GPU card itself. Put ! in front to run a shell command in Colab.

!nvidia-smi

Option 2 — Local Install with pip

Linux and macOS

python3 -m venv tf-env
source tf-env/bin/activate
pip install --upgrade pip
pip install tensorflow              # CPU (and Apple Silicon Macs)
# pip install 'tensorflow[and-cuda]'  # Linux + NVIDIA GPU

Windows (PowerShell)

py -m venv tf-env
tf-env\Scripts\activate
python -m pip install --upgrade pip
pip install tensorflow              # CPU only on native Windows

Verify the Install

import tensorflow as tf

print(tf.__version__)
print(tf.reduce_sum(tf.random.normal([1000, 1000])).dtype)
print("Works!")

Try It — Install Command Builder

How TensorFlow Talks to the GPU

Your Python code does not talk to the GPU directly. It passes through several layers. If any layer is missing, TensorFlow quietly falls back to the CPU.

GPU Check — The Code

Run this block on any machine. It tells you what TensorFlow can see.

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
print("Num GPUs available:", len(gpus))
for gpu in gpus:
    print("  ", gpu.name, tf.config.experimental.get_device_details(gpu).get('device_name'))

info = tf.sysconfig.get_build_info()
print("Built with CUDA:", info.get('is_cuda_build'))
print("CUDA version   :", info.get('cuda_version'))
print("cuDNN version  :", info.get('cudnn_version'))

Where Did My Op Run?

Every tensor remembers its device. You can also force a device with tf.device.

a = tf.random.normal([3, 3])
print("a lives on:", a.device)

with tf.device('/CPU:0'):
    b = tf.matmul(a, a)
print("b lives on:", b.device)

On a GPU machine, a.device ends with GPU:0.

Stop TensorFlow from Taking All GPU Memory

By default TensorFlow grabs almost all GPU memory at start. This blocks other notebooks. Turn on memory growth at the very top of your script.

gpus = tf.config.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)   # must run before any op

Feel the Speed — CPU vs GPU

import time

def bench(device, n=4000):
    with tf.device(device):
        x = tf.random.normal([n, n])
        tf.matmul(x, x)                       # warm-up run
        start = time.time()
        y = tf.matmul(x, x)
        _ = y.numpy()                         # wait until finished
    return time.time() - start

print(f"CPU: {bench('/CPU:0'):.3f} s")
if tf.config.list_physical_devices('GPU'):
    print(f"GPU: {bench('/GPU:0'):.3f} s")

Try It — Setup Troubleshooter

Golden Rules