add inplace argument

2noise · Aug 10, 2024 · 677722a · 677722a
1 parent a7fb442
commit 677722a
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/ChatTTS/model/speaker.py b/ChatTTS/model/speaker.py
@@ -26,6 +26,7 @@ def apply(
         input_ids: torch.Tensor,
         spk_emb_ids: int,
         device: torch.device,
+        inplace: bool = True,
     ):
         if isinstance(spk_emb, str):
             spk_emb_tensor = torch.from_numpy(self._decode(spk_emb))
@@ -45,8 +46,10 @@ def apply(
             .expand(emb.shape)
         )
         cond = input_ids.narrow(-1, 0, 1).eq(spk_emb_ids).expand(emb.shape)
-        torch.where(cond, n, emb, out=emb)
-        del cond, n
+        out = torch.where(cond, n, emb, out=emb if inplace else None)
+        if inplace:
+            del cond, n
+        return out
 
     @staticmethod
     @torch.no_grad()