Update yuan_hf_model.py

yuan_hf_model.py

@@ -32,8 +32,8 @@ from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
 from .configuration_yuan import YuanConfig
 from einops import rearrange
-#from flash_attn import flash_attn_varlen_func as flash_attn_unpadded_func
-#from flash_attn import flash_attn_func
+from flash_attn import flash_attn_varlen_func as flash_attn_unpadded_func
+from flash_attn import flash_attn_func
 
 import copy
 
@@ -268,8 +268,8 @@ class YuanAttention(nn.Module):
         is_first_step = False
         if use_cache:
             if past_key_value is None:
-                # inference_hidden_states_memory = torch.empty(bsz, 2, hidden_states.shape[2], dtype=hidden_states.dtype ,device=torch.cuda.current_device())
-                inference_hidden_states_memory = torch.empty(bsz, 2, hidden_states.shape[2], dtype=hidden_states.dtype)
+                inference_hidden_states_memory = torch.empty(bsz, 2, hidden_states.shape[2], dtype=hidden_states.dtype ,device=torch.cuda.current_device())
+                #inference_hidden_states_memory = torch.empty(bsz, 2, hidden_states.shape[2], dtype=hidden_states.dtype)
                 is_first_step = True
             else:
                 before_hidden_states = past_key_value[2]