Fix apply_repetition_penalty

streaming_mossttsrealtime.py

@@ -365,24 +365,14 @@ class MossTTSRealtimeInference:
         repetition_window: Optional[int] = None,
     ):
         scores_ = scores[:, 0, :]
-        batch_size = scores_.shape[0]
         ht = history_tokens
 
         if repetition_window is not None and repetition_window > 0:
             ht = ht[:, -repetition_window:]
 
-        ht_sorted, _ = torch.sort(ht, dim=1)
-        uniq = torch.unique_consecutive(ht_sorted, dim=1)
-
-        b_idx = torch.arange(batch_size, device=uniq.device).unsqueeze(1).expand_as(uniq)
-        b_flat = b_idx.reshape(-1)
-        t_flat = uniq.reshape(-1)
-
-        cur = scores_[b_flat, t_flat]
+        cur = scores_.gather(1, ht)
         new = torch.where(cur < 0, cur * penalty, cur / penalty)
-
-        scores_[b_flat, t_flat] = new
-
+        scores_.scatter_(1, ht, new)
         return scores_
 
     def sample_token(self, logits, temperature, top_p=0.6, top_k=30, do_sample=True):