update monkey patch for llama2 (lm-sys#2379)

wangzhen263 · Sep 7, 2023 · 1d703b2 · 1d703b2
1 parent a5e6abf
commit 1d703b2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/fastchat/train/train_mem.py b/fastchat/train/train_mem.py
@@ -1,7 +1,7 @@
 # Make it more memory efficient by monkey patching the LLaMA model with FlashAttn.
 
 # Need to call this before importing transformers.
-from fastchat.train.llama_flash_attn_monkey_patch import (
+from fastchat.train.llama2_flash_attn_monkey_patch import (
     replace_llama_attn_with_flash_attn,
 )