codec/models/pardi_tokenizer.py

import torch

from zcodec.models import WavVAE, ZFlowAutoEncoder
from zcodec.models.wavvae.model import WavVAEConfig
from zcodec.models.zflowae.model import ZFlowAutoEncoderConfig


class PardiTokenizer(nn.Module):
    def __init__(self, wavvae_cfg: WavVAEConfig, zflowae_cfg: ZFlowAutoEncoderConfig):