load jit to device

FunAudioLLM · Sep 26, 2024 · ba3d969 · ba3d969
1 parent 06934c3
commit ba3d969
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/cosyvoice/cli/model.py b/cosyvoice/cli/model.py
@@ -63,11 +63,11 @@ def load(self, llm_model, flow_model, hift_model):
         self.hift.to(self.device).eval()
 
     def load_jit(self, llm_text_encoder_model, llm_llm_model, flow_encoder_model):
-        llm_text_encoder = torch.jit.load(llm_text_encoder_model)
+        llm_text_encoder = torch.jit.load(llm_text_encoder_model, map_location=self.device)
         self.llm.text_encoder = llm_text_encoder
-        llm_llm = torch.jit.load(llm_llm_model)
+        llm_llm = torch.jit.load(llm_llm_model, map_location=self.device)
         self.llm.llm = llm_llm
-        flow_encoder = torch.jit.load(flow_encoder_model)
+        flow_encoder = torch.jit.load(flow_encoder_model, map_location=self.device)
         self.flow.encoder = flow_encoder
 
     def load_onnx(self, flow_decoder_estimator_model):