Make using GPU in relaxation optional, with default not to use.

PiperOrigin-RevId: 471264036
Change-Id: Ic2201048f5c9d9d741cb20fdf46f2f27986587f4

notebooks/AlphaFold.ipynb

@@ -478,6 +478,12 @@
         "\n",
         "run_relax = True  #@param {type:\"boolean\"}\n",
         "\n",
+        "#@markdown Relaxation is faster with a GPU, but we have found it to be less stable.\n",
+        "#@markdown You may wish to enable GPU for higher performance, but if it doesn't\n",
+        "#@markdown converge we suggested reverting to using without GPU.\n",
+        "\n",
+        "relax_use_gpu = False  #@param {type:\"boolean\"}\n",
+        "\n",
         "# --- Run the model ---\n",
         "if model_type_to_use == notebook_utils.ModelType.MONOMER:\n",
         "  model_names = config.MODEL_PRESETS['monomer'] + ('model_2_ptm',)\n",
@@ -554,7 +560,7 @@
         "        stiffness=10.0,\n",
         "        exclude_residues=[],\n",
         "        max_outer_iterations=3,\n",
-        "        use_gpu=True)\n",
+        "        use_gpu=relax_use_gpu)\n",
         "    relaxed_pdb, _, _ = amber_relaxer.process(prot=unrelaxed_proteins[best_model_name])\n",
         "  else:\n",
         "    print('Warning: Running without the relaxation stage.')\n",