fix bug for mefture (#5299)

10 months ago · ddf879e2db
parent f7e3f82a7e
commit ddf879e2db
1 changed files with 4 additions and 1 deletions
--- 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):