From 447587300b75b71f10835face960a27324f4aefc Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Wed, 29 Apr 2026 09:16:50 -0700 Subject: [PATCH] fix: include medusa in data_module assignment in main.py When mode == "medusa", data_module was never assigned because the condition only covered "eagle3" and "dflash", causing an UnboundLocalError at the trainer construction. Add "medusa" to the condition so the data module is correctly prepared for all supported training modes. Fixes OMNIML-4147 Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Ye Yu --- examples/speculative_decoding/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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,