From 28fecc2f7665bbc5540086bac218492fd1b180eb Mon Sep 17 00:00:00 2001 From: Oh Joon Kwon Date: Sat, 14 Dec 2024 00:16:13 +0900 Subject: [PATCH] Keeping it simple. Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com> --- src/accelerate/utils/other.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/accelerate/utils/other.py b/src/accelerate/utils/other.py index ce358a6acdc..5c1c5b23dbb 100644 --- a/src/accelerate/utils/other.py +++ b/src/accelerate/utils/other.py @@ -128,10 +128,9 @@ def _recursive_unwrap(module): if getattr(model, "_converted_to_transformer_engine", False): convert_model(model, to_transformer_engine=False) - if keep_torch_compile: - if is_compiled: - compiled_model._orig_mod = model - model = compiled_model + if keep_torch_compile and is_compiled: + compiled_model._orig_mod = model + model = compiled_model return model