Improve code

lenML · Jun 6, 2024 · 0e278ab · 0e278ab
1 parent 9764e23
commit 0e278ab
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/modules/ChatTTS/ChatTTS/model/dvae.py b/modules/ChatTTS/ChatTTS/model/dvae.py
@@ -143,9 +143,9 @@ def forward(self, inp):
         else:
             vq_feats = inp.detach().clone()
 
-        temp = torch.chunk(vq_feats, 2, dim=1) # flatten trick :)
-        temp = torch.stack(temp, -1)
-        vq_feats = temp.reshape(*temp.shape[:2], -1)
+        vq_feats = vq_feats.view(
+            (vq_feats.size(0), 2, vq_feats.size(1)//2, vq_feats.size(2)),
+        ).permute(0, 2, 3, 1).flatten(2)
 
         vq_feats = vq_feats.transpose(1, 2)
         dec_out = self.decoder(input=vq_feats)