torchrun \
--nnodes=1 \
--nproc_per_node=8 \
train.py \
--model SiT-S/2 \
--epochs=800 \
--data-path /home/nvidia/datasets/imagenet-1k/train \
--wandb \
--global-batch-size=1024 \
--run-name 005-SiT-S-2-Linear-velocity-None \
--ckpt /home/nvidia/SiT-Complementary/results/005-SiT-S-2-Linear-velocity-None/checkpoints/0550000.pt

# batch_size x 4, lr x 2