Add fallback to draft_rope_scale to 1.0

lonestriker · Dec 6, 2023 · 3f2fcbc · 3f2fcbc
1 parent 39f7a2a
commit 3f2fcbc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/model.py b/model.py
@@ -111,7 +111,7 @@ def progress(loaded_modules: int, total_modules: int, loading_draft: bool)
             self.draft_config.model_dir = str(draft_model_path.resolve())
             self.draft_config.prepare()
 
-            if "draft_rope_scale" in kwargs: self.draft_config.scale_pos_emb = kwargs["draft_rope_scale"]
+            self.draft_config.scale_pos_emb = kwargs.get("draft_rope_scale") or 1.0
             self.draft_config.scale_alpha_value = kwargs.get("draft_rope_alpha") or self.calculate_rope_alpha(self.draft_config.max_seq_len)
             self.draft_config.max_seq_len = self.config.max_seq_len