Load schedulers from ckpt in a backward-compatible manner (#2236)

This commit is contained in:
hyperzy authored and GitHub committed 2023-07-18 18:02:15 +00:00
1 parent 9f487a464f
commit e7fd1aefa2
2 files changed
+6 -1

No files matched your search

+4 -1
View File
@@ -406,7 +406,8 @@ class Trainer:
# load the checkpoints for pipeline, optimizers, and gradient scalar
self.pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"])
self.optimizers.load_optimizers(loaded_state["optimizers"])
self.optimizers.load_schedulers(loaded_state["schedulers"])
if "schedulers" in loaded_state and self.config.load_scheduler:
self.optimizers.load_schedulers(loaded_state["schedulers"])
self.grad_scaler.load_state_dict(loaded_state["scalers"])
CONSOLE.print(f"Done loading Nerfstudio checkpoint from {load_path}")
elif load_checkpoint is not None:
@@ -416,6 +417,8 @@ class Trainer:
# load the checkpoints for pipeline, optimizers, and gradient scalar
self.pipeline.load_pipeline(loaded_state["pipeline"], loaded_state["step"])
self.optimizers.load_optimizers(loaded_state["optimizers"])
if "schedulers" in loaded_state and self.config.load_scheduler:
self.optimizers.load_schedulers(loaded_state["schedulers"])
self.grad_scaler.load_state_dict(loaded_state["scalers"])
CONSOLE.print(f"Done loading Nerfstudio checkpoint from {load_checkpoint}")
else: