Keeping it simple.

Co-authored-by: Marc Sun <[email protected]>
huggingface · Dec 13, 2024 · 28fecc2 · 28fecc2
1 parent c3e69c5
commit 28fecc2
Showing 1 changed file with 3 additions and 4 deletions.
diff --git 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