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.
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:
If this prints False, you are either on a CPU runtime or the GPU is not currently attached to the session.
With TensorFlow:
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.
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:
That produces compact output such as:
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:
TensorFlow:
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-smiprogrammatically 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
- how to get covariance matrix in tensorflow?
- How to get current available GPUs in tensorflow?
- How to get intermediate layers' output of pre-trained BERT model in HuggingFace Transformers library?
- How to get labels ids in Keras when training on multiple classes?
- How to get both MSE and R2 from a sklearn GridSearchCV?
- How to get classification probabilities from PySpark MultilayerPerceptronClassifier?
- how to get covariance matrix in tensorflow?
- How to get different colored lines for different plots in a single figure
.png&w=3840&q=75)
Tackling System Design Interview Problems
A short course that equips you with the skills to approach system design interviews methodically.
Start the free courseTrack 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.