Fix error wrongly triggered on model key absence

facebookresearch · Dec 31, 2024 · 496fc36 · 496fc36
1 parent 03511c2
commit 496fc36
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/fairseq2/recipes/llama/convert_checkpoint.py b/src/fairseq2/recipes/llama/convert_checkpoint.py
@@ -122,8 +122,8 @@ def run(self, args: Namespace) -> int:
 
                     sys.exit(1)
 
-                if "model" not in checkpoint:
-                    log.error("Checkpoint file {} does not contain a 'model' entry.", input_file.name)  # fmt: skip
+                if all(key not in checkpoint for key in ["model_key", "model"]):
+                    log.error("Checkpoint file {} does not contain a 'model_key' nor 'model' entry.", input_file.name)  # fmt: skip
 
                     sys.exit(1)