fix device mismatch bug

2noise · Aug 31, 2024 · 1d5f1fc · 1d5f1fc
1 parent 024f93e
commit 1d5f1fc
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/ChatTTS/core.py b/ChatTTS/core.py
@@ -288,8 +288,8 @@ def _load(
             self.config.embed.num_text_tokens,
             self.config.embed.num_vq,
         )
-        embed.from_pretrained(embed_path)
-        self.embed = embed
+        embed.from_pretrained(embed_path, device=self.device)
+        self.embed = embed.to(self.device)
         self.logger.log(logging.INFO, "embed loaded.")
 
         gpt = GPT(

diff --git a/ChatTTS/model/embed.py b/ChatTTS/model/embed.py
@@ -34,12 +34,13 @@ def __init__(
         )
 
     @torch.inference_mode()
-    def from_pretrained(self, filename: str):
+    def from_pretrained(self, filename: str, device: torch.device):
         state_dict_tensors = {}
         with safe_open(filename, framework="pt") as f:
             for k in f.keys():
                 state_dict_tensors[k] = f.get_tensor(k)
         self.load_state_dict(state_dict_tensors)
+        self.to(device)
 
     def __call__(
         self, input_ids: torch.Tensor, text_mask: torch.Tensor