Set environment variables before any other code is executed.

PiperOrigin-RevId: 495028896
Change-Id: I3b1a6ddceca1961bffdaba443e7db47bbfbc4fad

notebooks/AlphaFold.ipynb

@@ -65,6 +65,11 @@
       },
       "outputs": [],
       "source": [
+        "# Set environment variables before running any other code.\n",
+        "import os\n",
+        "os.environ['TF_FORCE_UNIFIED_MEMORY'] = '1'\n",
+        "os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '4.0'\n",
+        "\n",
         "#@title 1. Install third-party software\n",
         "\n",
         "#@markdown Please execute this cell by pressing the _Play_ button\n",
@@ -190,12 +195,7 @@
         "# Make sure everything we need is on the path.\n",
         "import sys\n",
         "sys.path.append('/opt/conda/lib/python3.8/site-packages')\n",
-        "sys.path.append('/content/alphafold')\n",
-        "\n",
-        "# Make sure all necessary environment variables are set.\n",
-        "import os\n",
-        "os.environ['TF_FORCE_UNIFIED_MEMORY'] = '1'\n",
-        "os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '2.0'"
+        "sys.path.append('/content/alphafold')"
       ]
     },
     {