CosyVoice commited on
Commit
ba3d969
1 Parent(s): 06934c3

load jit to device

Browse files
Files changed (1) hide show
  1. cosyvoice/cli/model.py +3 -3
cosyvoice/cli/model.py CHANGED
@@ -63,11 +63,11 @@ class CosyVoiceModel:
63
  self.hift.to(self.device).eval()
64
 
65
  def load_jit(self, llm_text_encoder_model, llm_llm_model, flow_encoder_model):
66
- llm_text_encoder = torch.jit.load(llm_text_encoder_model)
67
  self.llm.text_encoder = llm_text_encoder
68
- llm_llm = torch.jit.load(llm_llm_model)
69
  self.llm.llm = llm_llm
70
- flow_encoder = torch.jit.load(flow_encoder_model)
71
  self.flow.encoder = flow_encoder
72
 
73
  def load_onnx(self, flow_decoder_estimator_model):
 
63
  self.hift.to(self.device).eval()
64
 
65
  def load_jit(self, llm_text_encoder_model, llm_llm_model, flow_encoder_model):
66
+ llm_text_encoder = torch.jit.load(llm_text_encoder_model, map_location=self.device)
67
  self.llm.text_encoder = llm_text_encoder
68
+ llm_llm = torch.jit.load(llm_llm_model, map_location=self.device)
69
  self.llm.llm = llm_llm
70
+ flow_encoder = torch.jit.load(flow_encoder_model, map_location=self.device)
71
  self.flow.encoder = flow_encoder
72
 
73
  def load_onnx(self, flow_decoder_estimator_model):