diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py
index 4a5f6288854..7e399cf9603 100644
--- a/examples/speculative_decoding/main.py
+++ b/examples/speculative_decoding/main.py
@@ -341,7 +341,7 @@ def train():
 
     print_rank_0("Loading dataset...")
     is_dflash = training_args.mode == "dflash"
-    if training_args.mode in ("eagle3", "dflash"):
+    if training_args.mode in ("eagle3", "medusa", "dflash"):
         data_module = make_speculative_data_module(
             tokenizer,
             data_args,