diff --git a/machine-learning/xgboost.ipynb b/machine-learning/xgboost.ipynb
index 22a56392..e8f9247e 100644
--- a/machine-learning/xgboost.ipynb
+++ b/machine-learning/xgboost.ipynb
@@ -114,7 +114,7 @@
    "source": [
     "import dask\n",
     "import xgboost\n",
-    "import dask_xgboost"
+    "from xgboost.dask import DaskDMatrix, train"
    ]
   },
   {
@@ -140,8 +140,10 @@
     "params = {'objective': 'binary:logistic',\n",
     "          'max_depth': 4, 'eta': 0.01, 'subsample': 0.5, \n",
     "          'min_child_weight': 0.5}\n",
+    "dtrain = DaskDMatrix(client, X_train, y_train)\n",
     "\n",
-    "bst = dask_xgboost.train(client, params, X_train, y_train, num_boost_round=10)"
+    "output = train(client, params, dtrain, num_boost_round=10)\n",
+    "bst = output['booster']  # This is the actual Booster object\n"
    ]
   },
   {
@@ -219,7 +221,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "y_hat = dask_xgboost.predict(client, bst, X_test).persist()\n",
+    "from xgboost.dask import predict\n",
+    "y_hat = predict(client, bst, X_test).persist()\n",
     "y_hat"
    ]
   },
@@ -295,7 +298,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.12"
+   "version": "3.12.3"
   }
  },
  "nbformat": 4,