diff --git a/.github/workflows/examples-ci.yaml b/.github/workflows/examples-ci.yaml
index bacb651..ddc1035 100644
--- a/.github/workflows/examples-ci.yaml
+++ b/.github/workflows/examples-ci.yaml
@@ -30,7 +30,7 @@ jobs:
     strategy:
         fail-fast: false
         matrix:
-          os: [macOS-12, ubuntu-latest]
+          os: [ubuntu-latest]
           python-version: ["3.11", "3.12"]
           pydantic-version: ["2"]
           include-rdkit: [true]
@@ -89,4 +89,4 @@ jobs:
 
     - name: Run example notebooks
       run: |
-        python -m pytest -r fE -v -x --tb=short -nauto --nbval-lax --dist loadscope examples
+        python -m pytest -r fE -v -x --tb=short -nauto --nbval-lax --nbval-cell-timeout=50000 --dist loadscope examples
diff --git a/examples/train-gnn-to-electric-field/train-gnn-to-electric-field.ipynb b/examples/train-gnn-to-electric-field/train-gnn-to-electric-field.ipynb
index 0defa7f..2e918cd 100644
--- a/examples/train-gnn-to-electric-field/train-gnn-to-electric-field.ipynb
+++ b/examples/train-gnn-to-electric-field/train-gnn-to-electric-field.ipynb
@@ -569,6 +569,30 @@
     "from pytorch_lightning.callbacks import TQDMProgressBar"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c8cb0a08-eea6-41e4-9fd6-158204e26e79",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = pl.Trainer(\n",
+    "    max_epochs=100,\n",
+    "    callbacks=[TQDMProgressBar()], # add progress bar\n",
+    "    accelerator=\"cpu\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f8846142-77f2-4d60-a74e-c3bc3c3ed3ee",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "datamodule = training_model.create_data_module(verbose=False)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 14,
@@ -2058,16 +2082,9 @@
     }
    ],
    "source": [
-    "trainer = pl.Trainer(\n",
-    "    max_epochs=100,\n",
-    "    callbacks=[TQDMProgressBar()], # add progress bar\n",
-    "    accelerator=\"cpu\"\n",
-    ")\n",
-    "\n",
-    "\n",
     "trainer.fit(\n",
     "    training_model,\n",
-    "    datamodule=training_model.create_data_module(verbose=False),\n",
+    "    datamodule=datamodule,\n",
     ")"
    ]
   },
@@ -2137,7 +2154,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.5"
+   "version": "3.11.4"
   }
  },
  "nbformat": 4,