diff --git a/example.py b/example.py
index 2ce0b6a..77151d2 100644
--- a/example.py
+++ b/example.py
@@ -8,19 +8,15 @@ model = AIIAmoe(config)
 model.save_pretrained("test")
 model = AIIAmoe.from_pretrained("test")
 
-# Initialize pretrainer with the model
 pretrainer = Pretrainer(model, learning_rate=1e-4, config=config)
 
-# List of dataset paths
-dataset_paths = [
-    "/path/to/dataset1.parquet",
-    "/path/to/dataset2.parquet"
-]
+# Set checkpoint directory
+checkpoint_dir = "checkpoints/my_model"
 
-# Start training with multiple datasets
+# Start training (will automatically load checkpoint if available)
 pretrainer.train(
-    dataset_paths=dataset_paths,
-    num_epochs=10,
-    batch_size=2,
-    sample_size=10000
-)
\ No newline at end of file
+    dataset_paths=["path/to/dataset1.parquet", "path/to/dataset2.parquet"],
+    output_path="trained_models/my_model",
+    checkpoint_dir=checkpoint_dir,
+    num_epochs=10
+)