    tokens_config = get_start_tokens(config["quantizer"], n_tokens)
    quantizer = AudioTokenizer(config["quantizer"], tokens_config)
    codebook_size = config["quantizer"]["speech"]["n_new_tokens"] + config["quantizer"]["wav"]["n_new_tokens"]
    train_dataset, val_dataset = load_data(data, tokenizer, quantizer, config)
    model.resize_token_embeddings(n_tokens + codebook_size)

        model = fix_checkpoint(model, checkpoint_path)
    train_dataloader = DataLoader(
        train_dataset,
        shuffle=True,
        collate_fn=default_data_collator,
