    def forward(self, input_ids):
        """Input: [num_quant, B, T] -> Output: [B, T, H]"""
        num_quant, B, T = input_ids.shape
        summed_embeddings = torch.zeros(
            B, T, self.embeddings[0].embedding_dim, device=input_ids.device
        )
        for i in range(num_quant):
            summed_embeddings += self.embeddings[i](input_ids[i])
        return summed_embeddings
Input: [num_quant, B, T] -> Output: [B, T, H]
models/tts/valle_v2/valle_nar.py