Merge pull request #523 from google:doc1

PiperOrigin-RevId: 562575129
google · Sep 4, 2023 · 5bf8d3d · 5bf8d3d
2 parents 9ccf7d9 + 3d602e2
commit 5bf8d3d
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/docs/notebooks/distributed/custom_loop_pjit_example.md b/docs/notebooks/distributed/custom_loop_pjit_example.md
@@ -5,7 +5,7 @@ jupytext:
     extension: .md
     format_name: myst
     format_version: 0.13
-    jupytext_version: 1.14.4
+    jupytext_version: 1.15.1
 kernelspec:
   display_name: Python 3
   language: python
@@ -116,7 +116,7 @@ from typing import Any, Callable, Tuple, Union
 try:
     import jax.tools.colab_tpu
     jax.tools.colab_tpu.setup_tpu()
-except KeyError:
+except (KeyError, RuntimeError):
     print("TPU not found, continuing without it.")
 
 import jax

diff --git a/docs/notebooks/distributed/custom_loop_pmap_example.md b/docs/notebooks/distributed/custom_loop_pmap_example.md
@@ -5,7 +5,7 @@ jupytext:
     extension: .md
     format_name: myst
     format_version: 0.13
-    jupytext_version: 1.14.4
+    jupytext_version: 1.15.1
 kernelspec:
   display_name: Python 3
   language: python