add load_strict_model_weights as an optional config parameter (#655)

mosaicml · Oct 18, 2023 · 2c5965e · 2c5965e
1 parent cc238a3
commit 2c5965e
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/scripts/train/train.py b/scripts/train/train.py
@@ -383,6 +383,10 @@ def main(cfg: DictConfig) -> Trainer:
                                          'load_weights_only',
                                          must_exist=False,
                                          default_value=False)
+    load_strict_model_weights: bool = pop_config(cfg,
+                                                 'load_strict_model_weights',
+                                                 must_exist=False,
+                                                 default_value=True)
     load_ignore_keys: Optional[List[str]] = pop_config(cfg,
                                                        'load_ignore_keys',
                                                        must_exist=False,
@@ -567,6 +571,7 @@ def main(cfg: DictConfig) -> Trainer:
         save_weights_only=save_weights_only,
         load_path=load_path,
         load_weights_only=load_weights_only,
+        load_strict_model_weights=load_strict_model_weights,
         load_ignore_keys=load_ignore_keys,
         autoresume=autoresume,
         python_log_level=python_log_level,