    def forward_index(self, x, flatten_idx=False):
        """
        all_indices: [num_of_quantizers, B, T]
        """
        quantized_out = 0.0
        residual = x
        all_indices = []
        for i, layer in enumerate(self.layers):
            quantized, indices = layer.forward_index(residual)
            # residual = residual - quantized.detach()
            residual = residual - quantized
            quantized_out = quantized_out + quantized
            if flatten_idx:
                indices += self.codebook_size * i
            all_indices.append(indices)
        all_indices = torch.stack(all_indices)
        return quantized_out, all_indices
all_indices: [num_of_quantizers, B, T]
models/codec/vevo/vevo_repcodec.py