def __call__(self, batch: Dict[str, Any]) -> Dict[str, Any]:
        """
        Performs a forward pass through the model.

        Args:
            batch (dict): A dictionary containing features, reference waveform, and target waveform.

        Returns:
            dict: A dictionary containing the reconstruction, features, and other metrics.
        """
        feat = batch["feat"]
        ref_wav = batch["ref_wav"]
        mel = self.get_mel_spectrogram(ref_wav)
        z = self.encoder(feat.transpose(0, 2, 1))
        vq_outputs = self.quantizer(z)

        x_vector, d_vector = self.speaker_encoder(mel)

        conditions = d_vector
        with_speaker_loss = False

        # Ensure conditions is an integer type for embedding lookup
        # The error shows that the embedding layer expects integral indices
        if isinstance(conditions, mx.array) and conditions.dtype == mx.float32:
            # Convert to integer type if needed for the embedding layer
            # or ensure it's properly formatted for the prenet
            conditions = conditions.astype(mx.int32)

        x = self.prenet(vq_outputs["z_q"], conditions)
        pred_feat = self.postnet(x)
        x = x + conditions[..., None]
        wav_recon = self.decoder(x)

        return {
            "vq_loss": vq_outputs["vq_loss"],
            "perplexity": vq_outputs["perplexity"],
            "cluster_size": vq_outputs["active_num"],
            "recons": wav_recon,
            "pred_feat": pred_feat,
            "x_vector": x_vector,
            "d_vector": d_vector,
            "audios": batch["wav"][:, None],
            "with_speaker_loss": with_speaker_loss,
        }

        Performs a forward pass through the model.

        Args:
            batch (dict): A dictionary containing features, reference waveform, and target waveform.

        Returns:
            dict: A dictionary containing the reconstruction, features, and other metrics.
        
mlx_audio/tts/models/spark/bicodec.py