Skip missing keys in --checkpoint-keys

SimonLarsen · Dec 9, 2024 · ef12ca4 · ef12ca4
1 parent 70cdb36
commit ef12ca4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/frogbox/pipelines/pipeline.py b/frogbox/pipelines/pipeline.py
@@ -169,7 +169,7 @@ def _load_checkpoint(
     ) -> None:
         if keys is None:
             keys = list(to_load.keys())
-        to_load = {k: to_load[k] for k in keys}
+        to_load = {k: to_load[k] for k in keys if k in to_load}
 
         Checkpoint.load_objects(
             to_load=to_load,