Set os.environ["TF_USE_LEGACY_KERAS"] = "1" in Colab notebooks

to guarantee a compatible version of `tf.keras`. PiperOrigin-RevId: 604380766
tensorflow · Feb 6, 2024 · 049d41c · 049d41c
1 parent ffa453f
commit 049d41c
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 0 deletions.
diff --git a/examples/notebooks/intro_mutag_example.ipynb b/examples/notebooks/intro_mutag_example.ipynb
@@ -116,6 +116,7 @@
         "\n",
         "import matplotlib.pyplot as plt\n",
         "import numpy as np\n",
+        "os.environ[\"TF_USE_LEGACY_KERAS\"] = \"1\"  # Select tf.keras version 2, not 3.\n",
         "import tensorflow as tf\n",
         "import tensorflow_gnn as tfgnn\n",
         "\n",

diff --git a/examples/notebooks/ogbn_mag_e2e.ipynb b/examples/notebooks/ogbn_mag_e2e.ipynb
@@ -120,6 +120,7 @@
         "import re\n",
         "from typing import Mapping\n",
         "\n",
+        "os.environ[\"TF_USE_LEGACY_KERAS\"] = \"1\"  # Select tf.keras version 2, not 3.\n",
         "import tensorflow as tf\n",
         "import tensorflow_gnn as tfgnn\n",
         "from tensorflow_gnn import runner\n",

diff --git a/examples/notebooks/ogbn_mag_indepth.ipynb b/examples/notebooks/ogbn_mag_indepth.ipynb
@@ -128,6 +128,7 @@
         "import re\n",
         "\n",
         "from google.protobuf import text_format\n",
+        "os.environ[\"TF_USE_LEGACY_KERAS\"] = \"1\"  # Select tf.keras version 2, not 3.\n",
         "import tensorflow as tf\n",
         "import tensorflow_gnn as tfgnn\n",
         "\n",