diff --git a/python/flexflow/serve/serve.py b/python/flexflow/serve/serve.py
index cbc4122897..ac622b3337 100644
--- a/python/flexflow/serve/serve.py
+++ b/python/flexflow/serve/serve.py
@@ -373,9 +373,9 @@ def compile(
         model_configs = self.config_class(self.hf_config)
 
         self.rm.set_max_spec_tree_token_num(
-            self.model_configs.max_spec_tree_token_num
+            model_configs.max_spec_tree_token_num
             if "max_spec_tree_token_num"
-            in self.model_configs.__dict__
+            in model_configs.__dict__
             else 20
         )