srun --label python -u deit/main.py \
        --model mega_base_patch16_224 \
        --batch-size 128 \
        --lr 2e-3 \
        --seed 42 \
        --norm-type 'layernorm' \
        --drop-path 0.3 \
        --epochs 300 \
        --warmup-epochs 20 \
        --weight-decay 0.05 \
        --warmup-lr 1e-5 \
        --clip-grad 1.0 \
        --opt-betas 0.9 0.98 \
        --world_size 8 \
        --reprob 0.25 \
        --repeated-aug 3 \
        --data-path ${DATA_PATH} \
        --output_dir ${MODEL_PATH}