diff --git a/RELEASE.md b/RELEASE.md
index 248b6461..09591ed8 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -1,5 +1,11 @@
 # Learning Interpretability Tool Release Notes
 
+
+## Release 1.3.1
+
+This is a minor update to fix issues with running the [LIT Gemma Colab](https://colab.sandbox.google.com/github/google/generative-ai-docs/blob/main/site/en/gemma/docs/lit_gemma.ipynb).
+
+
 ## Release 1.3
 
 This release updates how the Learning Interpretability Tool (LIT) can be
diff --git a/lit_nlp/notebook.py b/lit_nlp/notebook.py
index 02ad0e4e..b1bad3e6 100644
--- a/lit_nlp/notebook.py
+++ b/lit_nlp/notebook.py
@@ -22,22 +22,18 @@
 from lit_nlp import server_config
 from lit_nlp.api import layout
 from lit_nlp.lib import wsgi_serving
+from tqdm import notebook
 
 JsonDict = Mapping[str, Any]
 
-is_colab = False
 try:
   import google.colab  # pylint: disable=g-import-not-at-top,unused-import
   from google.colab import output  # pylint: disable=g-import-not-at-top,unused-import # pytype: disable=import-error
   is_colab = True
-  # Can disable import error as this package is always
-  # included in colab kernels.
-  from colabtools import interactive_widgets  # pylint: disable=g-import-not-at-top # pytype: disable=import-error
-  progress_indicator = interactive_widgets.ProgressIter
 except (ImportError, ModuleNotFoundError):
-  from tqdm import notebook  # pylint: disable=g-import-not-at-top
-  progress_indicator = notebook.tqdm
+  is_colab = False
 
+progress_indicator = notebook.tqdm
 modules = layout.LitModuleName
 
 LIT_NOTEBOOK_LAYOUT = layout.LitCanonicalLayout(
diff --git a/lit_nlp/package.json b/lit_nlp/package.json
index b914daf4..42a26bdd 100644
--- a/lit_nlp/package.json
+++ b/lit_nlp/package.json
@@ -1,6 +1,6 @@
 {
   "name": "lit-client",
-  "version": "1.3",
+  "version": "1.3.1",
   "description": "Client app for LIT, the Learning Interpretability Tool",
   "engines": {
     "node": ">=18.0.0"
diff --git a/pyproject.toml b/pyproject.toml
index 831d06e8..2a27db38 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 name = "lit-nlp"
 # LINT.IfChange
-version = "1.3"
+version = "1.3.1"
 # LINT.ThenChange(./lit_nlp/package.json)
 authors = [
   { name="Google, LLC", email="lit-dev@google.com" }
@@ -40,7 +40,6 @@ dependencies = [
   "six>=1.16.0",
   "termcolor>=2.3.0",
   "tqdm>=4.64.0",
-  "umap-learn>=0.5.1",
   "werkzeug>=2.2.3",
 ]
 # LINT.ThenChange(./requirements.txt)
@@ -105,10 +104,14 @@ examples-generative-ai = [
   "vertexai>=1.49.0",
 ]
 # LINT.ThenChange(./requirements_examples_generative_ai.txt)
+umap = [
+  "umap-learn==0.5.6"
+]
 # LINT.IfChange
 test = [
   "lit-nlp[examples-discriminative-ai]",
   "lit-nlp[examples-generative-ai]",
+  "lit-nlp[umap]",
   "lime==0.2.0.1",
   "pytest>=7.4.0,<8.0.0",
   "webtest>=2.0",
diff --git a/requirements.txt b/requirements.txt
index 6e740ce9..9c4707b0 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -37,6 +37,5 @@ shap>=0.42.0,<0.46.0
 six>=1.16.0
 termcolor>=2.3.0
 tqdm>=4.64.0
-umap-learn>=0.5.1
 werkzeug>=2.2.3
 # LINT.ThenChange(./pyproject.toml)
diff --git a/requirements_test.txt b/requirements_test.txt
index 367a238f..9e2e0556 100644
--- a/requirements_test.txt
+++ b/requirements_test.txt
@@ -18,5 +18,6 @@
 
 lime==0.2.0.1
 pytest>=7.4.0,<8.0.0
+umap-learn==0.5.6
 webtest>=2.0
 # LINT.ThenChange(./pyproject.toml)
diff --git a/website/sphinx_src/index.md b/website/sphinx_src/index.md
index 64921be0..205e830c 100644
--- a/website/sphinx_src/index.md
+++ b/website/sphinx_src/index.md
@@ -1,6 +1,6 @@
 # Learning Interpretability Tool (LIT)
 
-<!--* freshness: { owner: 'lit-dev' reviewed: '2024-08-15' } *-->
+<!--* freshness: { owner: 'lit-dev' reviewed: '2024-10-25' } *-->
 
 <!-- [TOC] placeholder - DO NOT REMOVE -->