Add the ValueError in the restore checkpoint exception to avoid problems with the optimizer restauration when new keys are addition

2026-09-09 11:16:00 +08:00 · 2021-12-09 13:27:21 +00:00
parent ecf327a118
commit fee01daa09
1 changed files with 1 additions and 1 deletions
@@ -453,7 +453,7 @@ class Trainer:
            if "scaler" in checkpoint and self.use_amp_scaler and checkpoint["scaler"]:
                print(" > Restoring Scaler...")
                scaler = _restore_list_objs(checkpoint["scaler"], scaler)
-        except (KeyError, RuntimeError):
+        except (KeyError, RuntimeError, ValueError):
            print(" > Partial model initialization...")
            model_dict = model.state_dict()
            model_dict = set_init_dict(model_dict, checkpoint["model"], config)