How to save & load ML models in Google Colab (Google Drive)

Colab runtimes reset. If you don’t save your trained model to Google Drive, you can lose hours of training. In this guide we’ll save and reload models for TensorFlow and PyTorch – and show where NoteCapsule fits in.

1. Mount Google Drive in Colab

from google.colab import drive
drive.mount("/content/drive")

Your Drive will be available at /content/drive/MyDrive/.

2. Save & load TensorFlow / Keras models

2.1 Save the entire model

model_save_path = "/content/drive/MyDrive/models/my_keras_model.keras"
model.save(model_save_path)

2.2 Load it later

import tensorflow as tf

loaded_model = tf.keras.models.load_model(model_save_path)
preds = loaded_model(test_data)

3. Save & load PyTorch models

3.1 Save state_dict to Drive

import torch
import os

MODEL_DIR = "/content/drive/MyDrive/models"
os.makedirs(MODEL_DIR, exist_ok=True)

model_save_path = os.path.join(MODEL_DIR, "classifier.pt")
torch.save(model.state_dict(), model_save_path)

3.2 Load in another notebook

model = MyModel()
model.load_state_dict(torch.load(model_save_path, map_location="cpu"))
model.eval()

Make sure the model class definition matches the one used when saving.

4. Keep a full project snapshot with NoteCapsule

Saving only the model weights can be risky if you forget which notebook, data, or preprocessing created them. NoteCapsule lets you save a fuller snapshot:

from notebookcapsule import create_capsule

create_capsule(
    name="after-best-model",
    notebook_path="notebooks/train.ipynb",
    data_dirs=["./data", "./models"],
    base_dir=".",  # project root
)

Now you have both the model weights and an organized Capsule describing the exact project state when you trained it.