diff --git a/examples/sdk/train_api.ipynb b/examples/sdk/train_api.ipynb
index 76c74a0354..242efe8ebb 100644
--- a/examples/sdk/train_api.ipynb
+++ b/examples/sdk/train_api.ipynb
@@ -73,9 +73,12 @@
     "        ),\n",
     "        training_parameters=TrainingArguments(\n",
     "            num_train_epochs=1,\n",
-    "            per_device_train_batch_size=4,\n",
-    "            gradient_accumulation_steps=4,\n",
+    "            per_device_train_batch_size=1,\n",
+    "            gradient_accumulation_steps=1,\n",
     "            gradient_checkpointing=True,\n",
+    "            gradient_checkpointing_kwargs={\n",
+    "                \"use_reentrant\": False\n",
+    "            },  # this is mandatory if checkpointng is enabled\n",
     "            warmup_steps=0.02,\n",
     "            learning_rate=1,\n",
     "            lr_scheduler_type=\"cosine\",\n",
@@ -93,7 +96,7 @@
     "    resources_per_worker={\n",
     "        \"gpu\": 1,\n",
     "        \"cpu\": 8,\n",
-    "        \"memory\": \"16Gi\",\n",
+    "        \"memory\": \"8Gi\",\n",
     "    },  # remove the gpu key if you don't want to attach gpus to the pods\n",
     ")"
    ]