Revert "[SW-199559] use native checkpointing under compile mode" (#1365)

huggingface · Sep 26, 2024 · 29d8604 · 29d8604
1 parent a4dbcff
commit 29d8604
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/optimum/habana/transformers/trainer.py b/optimum/habana/transformers/trainer.py
@@ -682,7 +682,7 @@ def _inner_training_loop(
 
             import transformers.modeling_utils
 
-            if args.deepspeed and args.use_lazy_mode:
+            if args.deepspeed:
                 from deepspeed.runtime.activation_checkpointing.checkpointing import (
                     CheckpointFunction,
                     non_reentrant_checkpoint,