echo "Starting train script..."
# export NCCL_DEBUG=INFO
# export NCCL_DEBUG_SUBSYS=ALL
# For multi-GPU training, use torchrun (modern way)
/home/minz/anaconda3/envs/nenv/bin/python -u train.py \
    --out_dir="/app2/suno/checkpoints" \
    --data_dir="/app2/suno/data/diffusion/v1" \
    --train_metas_filename="metas_diff_v4_val.jsonl" \
    --val_metas_filename="metas_diff_v4_val.jsonl" \
    \
    --encoder_dim=128 \
    --vae_dim=128 \
    --decoder_dim=1536 \
    --is_frozen_encoder=True \
    \
    --weight_mel_loss=15.0 \
    --weight_kl_loss=0.1 \
    --weight_feat_loss=2.0 \
    --weight_adv_loss=1.0 \
    --weight_disc_loss=1.0 \
    \
    --lr_scheduler_type="exponential" \
    --learning_rate_codec=1e-4 \
    --learning_rate_disc=1e-4 \
    --max_iters=500_000 \
    --eval_iters=20 \
    \
    --batch_size=55 \
    \
    --fsdp=False \
    --grad_checkpointing=True \
    --compile=False \
    --checkpoint_save_old_format=True \
    \
    --wandb_log=False \
    --wandb_dir="/app/suno/minz/wandb_logs" \
    --wandb_project="codec-2025q4" \
    --wandb_run_name="local_dac_finetune" \
    --local_cache_dir="/mnt/localdisk/tmp/minz"