From 1300054333cd48d49d997fc6748ff63160728669 Mon Sep 17 00:00:00 2001 From: Seungjun Nah Date: Wed, 23 Jul 2025 10:48:55 -0700 Subject: [PATCH] oom fix for conda environment when loading a checkpoint --- cosmos_predict2/checkpointer.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/cosmos_predict2/checkpointer.py b/cosmos_predict2/checkpointer.py index fa416996..cc9b628a 100644 --- a/cosmos_predict2/checkpointer.py +++ b/cosmos_predict2/checkpointer.py @@ -237,6 +237,7 @@ def load( strict=False if model.config.train_architecture == "lora" else True, ), ) + del state_dicts_to_load_for_dit_reg # Load EMA weights. if model.pipe.config.ema.enabled: set_model_state_dict( @@ -248,6 +249,7 @@ def load( strict=False if model.config.train_architecture == "lora" else True, ), ) + del state_dicts_to_load_for_dit_ema # Restore the attention operators. model.pipe.apply_cp()