update

inspiremusic/transformer/qwen_encoder.py

@@ -39,7 +39,7 @@ class QwenEncoder(nn.Module):
         else:
             self.dtype = torch.float32
 
-        self.model = AutoModelForCausalLM.from_pretrained(pretrain_path, device_map="auto", attn_implementation="flash_attention_2", torch_dtype=self.dtype)
+        self.model = AutoModelForCausalLM.from_pretrained(pretrain_path, device_map="cpu")
         self._output_size = self.model.config.hidden_size
         self.do_fusion_emb = do_fusion_emb
         self.hidden_norm = torch.nn.LayerNorm(self._output_size)
@@ -109,7 +109,7 @@ class QwenEmbeddingEncoder(nn.Module):
         else:
             self.dtype = torch.float32
         from transformers import Qwen2ForCausalLM
-        self.model = Qwen2ForCausalLM.from_pretrained(pretrain_path, device_map="auto", attn_implementation="flash_attention_2", torch_dtype=self.dtype)
+        self.model = Qwen2ForCausalLM.from_pretrained(pretrain_path, device_map="cpu")
         self._output_size = self.model.config.hidden_size
 
     def output_size(self) -> int: