model='DiT-g/2'
params=" \
            --qk-norm \
            --model ${model} \
            --rope-img base512 \
            --rope-real \
            "
export LD_LIBRARY_PATH=/public/tengcent-hy/HunyuanDiT-xiafy/rocblas-install-0530/lib:$LD_LIBRARY_PATH
export TORCHINDUCTOR_LAYOUT_OPTIMIZATION=1
export PYTORCH_MIOPEN_SUGGEST_NHWC=1
deepspeed hydit/train_deepspeed.py ${params}  "$@"