diff --git a/docs/mlflow.ipynb b/docs/mlflow.ipynb
index a8800d8..59b587c 100644
--- a/docs/mlflow.ipynb
+++ b/docs/mlflow.ipynb
@@ -50,8 +50,9 @@
    },
    "outputs": [],
    "source": [
-    "import lamindb as ln\n",
     "import lightning as pl\n",
+    "import lamindb as ln\n",
+    "from lamindb.integrations import lightning as ll\n",
     "import mlflow\n",
     "from pathlib import Path\n",
     "\n",
@@ -286,8 +287,8 @@
     "    autoencoder = LitAutoEncoder(hidden_size=32, bottleneck_size=16)\n",
     "\n",
     "    # Create a LaminDB LightningCallback which also (optionally) annotates checkpoints by desired metrics\n",
-    "    lamindb_callback = ln.integrations.lightning.Callback(\n",
-    "        path=Path(\"model_checkpoints\") / \"{mlflow_run.info.run_id}_last_epoch.ckpt\",\n",
+    "    lamindb_callback = ll.Callback(\n",
+    "        path=Path(\"model_checkpoints\") / f\"{mlflow_run.info.run_id}_last_epoch.ckpt\",\n",
     "        key=f\"testmodels/mlflow/{mlflow_run.info.run_id}.ckpt\",\n",
     "        features={\n",
     "            \"mlflow_run_id\": mlflow_run.info.run_id,\n",
diff --git a/docs/wandb.ipynb b/docs/wandb.ipynb
index 8c0d6b9..52f48df 100644
--- a/docs/wandb.ipynb
+++ b/docs/wandb.ipynb
@@ -46,9 +46,10 @@
    },
    "outputs": [],
    "source": [
+    "import lightning as pl\n",
     "import lamindb as ln\n",
+    "from lamindb.integrations import lightning as ll\n",
     "import wandb\n",
-    "import lightning as pl\n",
     "from pathlib import Path\n",
     "\n",
     "from torch import utils\n",
@@ -276,8 +277,8 @@
    "source": [
     "# Create a LaminDB LightningCallback which also (optionally) annotates checkpoints by desired metrics\n",
     "wandb_logger.experiment.id\n",
-    "lamindb_callback = ln.integrations.lightning.Callback(\n",
-    "    path=Path(\"model_checkpoints\") / \"{wanddblogger.version}_last_epoch.ckpt\",\n",
+    "lamindb_callback = ll.Callback(\n",
+    "    path=Path(\"model_checkpoints\") / f\"{wandb_logger.version}_last_epoch.ckpt\",\n",
     "    key=f\"testmodels/wandb/{wandb_logger.experiment.id}.ckpt\",\n",
     "    features={\n",
     "        \"wandb_run_id\": wandb_logger.experiment.id,\n",