from accelerate import Accelerator, DistributedDataParallelKwargs, InitProcessGroupKwargs
from src.data import load_data
from src.tokenizer import AudioTokenizer, get_start_tokens
from src.utils import save_checkpoint, fix_checkpoint, get_exp_name
# Parse arguments

parser.add_argument(
    "--config", type=str, help="Path to the config.yaml file", required=True
args = parser.parse_args()
# Load config
with open(args.config, "r") as file:
