fix bug for mefture (#5299)

ddf879e2 · Desperado-Jia · GitHub · f7e3f82a · ddf879e2
Unverified Commit ddf879e2 authored Jan 22, 2024 by Desperado-Jia Committed by GitHub Jan 22, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

applications/Colossal-LLaMA-2/colossal_llama2/utils/neftune_patch.py ...s/Colossal-LLaMA-2/colossal_llama2/utils/neftune_patch.py +4 -1

No files found.
--- a/applications/Colossal-LLaMA-2/colossal_llama2/utils/neftune_patch.py
+++ b/applications/Colossal-LLaMA-2/colossal_llama2/utils/neftune_patch.py
@@ -16,7 +16,10 @@ import torch
 def unwrap(model):
-    return model.unwrap().module
+    if hasattr(model, "module"):
+        return unwrap_model(model.module)
+    else:
+        return model
 def neftune_post_forward_hook(module, input, output):