Google Colab
GPU allocation
machine learning
Python
Jupyter Notebook

How to get allocated GPU spec in Google Colab

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

In Google Colab, the exact GPU you receive is assigned dynamically, so the first thing to do after connecting is inspect the runtime. The quickest answer is usually nvidia-smi, but you can also query the device programmatically from Python, TensorFlow, or PyTorch. That is useful when you want to log the hardware, choose a batch size, or decide whether mixed precision is worth enabling.

Confirm That the Runtime Has a GPU

First make sure the notebook is actually using a GPU runtime. In Colab, set the hardware accelerator to GPU and then check from Python:

python
import torch

print(torch.cuda.is_available())

If this prints False, you are either on a CPU runtime or the GPU is not currently attached to the session.

With TensorFlow:

python
import tensorflow as tf

print(tf.config.list_physical_devices("GPU"))

Either approach is enough to confirm that the runtime can see a GPU.

Use nvidia-smi for the Full Spec

The standard tool for seeing the assigned GPU model, memory, driver version, and utilization is nvidia-smi.

bash
!nvidia-smi

Typical output includes:

  • GPU model, such as Tesla T4 or A100
  • total memory
  • current memory usage
  • driver and CUDA versions
  • current utilization

For most users, this is the fastest and most complete answer.

Query the GPU Programmatically

If you want a clean string for logs or notebooks, query the values directly:

python
1import subprocess
2
3result = subprocess.check_output([
4    "nvidia-smi",
5    "--query-gpu=name,memory.total,driver_version",
6    "--format=csv,noheader"
7], text=True)
8
9print(result.strip())

That produces compact output such as:

text
Tesla T4, 15109 MiB, 535.129.03

This is useful when you want your notebook to record the environment automatically.

Framework-Specific Checks

Sometimes you only care about what your ML framework sees.

PyTorch:

python
1import torch
2
3if torch.cuda.is_available():
4    print(torch.cuda.get_device_name(0))
5    print(torch.cuda.get_device_properties(0).total_memory)

TensorFlow:

python
1import tensorflow as tf
2
3gpus = tf.config.list_physical_devices("GPU")
4print(gpus)

Framework checks are not as detailed as nvidia-smi, but they confirm that the runtime your code will use can actually access the device.

Interpret the Important Fields

When reading nvidia-smi, these fields matter most:

  • 'Name: the GPU model'
  • 'Memory-Usage: current used and total VRAM'
  • 'GPU-Util: current compute utilization'
  • 'CUDA Version: the driver-supported CUDA version'

If you are debugging out-of-memory errors, total memory is usually the critical field. If you are benchmarking training speed, the GPU model and runtime utilization matter more.

Understand Colab’s Constraints

Colab does not guarantee a specific GPU type. The same notebook may get a T4 one day and a different GPU another day depending on availability and account tier.

That means:

  • always inspect the runtime you actually received
  • avoid hard-coding expectations about model type
  • log the device in experiments so results are reproducible

This is one reason batch sizes or training times can vary between sessions.

Common Pitfalls

The most common mistake is assuming the notebook has a GPU because the code used one in a previous session. Colab runtimes are ephemeral, so check each time.

Another issue is reading nvidia-smi output before the runtime is fully initialized. If the command fails, reconnect or verify that GPU acceleration is actually enabled.

Some users also confuse driver-supported CUDA version with the exact CUDA toolkit environment available to Python packages. Those are related but not identical concepts.

Finally, do not assume the device name alone tells you everything about performance. Memory size, current load, and framework compatibility still matter.

Summary

  • In Colab, the quickest way to see the assigned GPU is !nvidia-smi.
  • Use PyTorch or TensorFlow checks to confirm your framework can access the GPU.
  • Query nvidia-smi programmatically when you want clean logging in notebooks.
  • GPU type in Colab is dynamic, so inspect the runtime each session.
  • The most important fields are usually model name, total memory, and current utilization.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.