diff --git a/examples/lightning_integration_examples/Part_2_lightning_advanced_and_gpu_tutorial.ipynb b/examples/lightning_integration_examples/Part_2_lightning_advanced_and_gpu_tutorial.ipynb
index 6755183e..147555f9 100644
--- a/examples/lightning_integration_examples/Part_2_lightning_advanced_and_gpu_tutorial.ipynb
+++ b/examples/lightning_integration_examples/Part_2_lightning_advanced_and_gpu_tutorial.ipynb
@@ -11,7 +11,9 @@
     "\n",
     "This notebook also describes how to use **multi-GPU** distributed training.\n",
     "\n",
-    "*For performance reasons, we only demonstrate single-GPU training. For a multi-GPU training example please refer to Part 6_lightning_multi_gpu.py*"
+    "*For performance reasons, we only demonstrate single-GPU training. For a multi-GPU training example please refer to Part 6_lightning_multi_gpu.py*\n",
+    "\n",
+    "**When running on Colab**: Please ensure your runtime has a GPU. Note that if using a single GPU runtime then the distributed GPU training example at the bottom of this notebook cannot be run"
    ]
   },
   {
@@ -276,11 +278,11 @@
    ],
    "source": [
     "# train with single GPU \n",
-    "lit_trainer = LitTrainer(epochs=200, accelerator='gpu', strategy='auto', devices=[1], custom_optimizer=optimizer)\n",
+    "lit_trainer = LitTrainer(epochs=200, accelerator='gpu', strategy='auto', devices=[0], custom_optimizer=optimizer)\n",
     "lit_trainer.fit(problem, data_setup_function)\n",
     "\n",
     "# train with two GPU. \n",
-    "lit_trainer = LitTrainer(epochs=200, accelerator='gpu', strategy='ddp_notebook', devices=[1,2], custom_optimizer=optimizer)\n",
+    "lit_trainer = LitTrainer(epochs=200, accelerator='gpu', strategy='ddp_notebook', devices=[0,1], custom_optimizer=optimizer)\n",
     "lit_trainer.fit(problem, data_setup_function)\n"
    ]
   },
diff --git a/pyproject.toml b/pyproject.toml
index 1154aa20..477369ae 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -53,7 +53,8 @@ dependencies = [
     "lightning", 
     "cvxpy", 
     "cvxpylayers", 
-    "casadi"
+    "casadi", 
+    "wandb"
 ]
 
 version = "1.5.0"
@@ -72,4 +73,4 @@ where = ["src"]
 [project.optional-dependencies]
 tests = ["pytest", "hypothesis"]
 examples = ["casadi", "cvxpy", "imageio", "cvxpylayers"]
-docs = ["sphinx", "sphinx-rtd-theme"]
\ No newline at end of file
+docs = ["sphinx", "sphinx-rtd-theme"]