Reproducing Issue on DiT-S/2

Hi, thanks for your great work!

I tried to reproduce the performance of DiT-S/2 on 256×256 resolution using this repository. My training settings followed both the configuration provided in the README and the default arguments in the training script, i.e.: • Global batch size = 256 • Learning rate = 1e-4

However, after 400k steps, the FID of the model was still far from the reported 68.4 in the paper. Even after 4000k steps, I was only able to achieve an FID of 79.59 on the test set.

For FID evaluation, I followed the paper’s settings: • cfg=1.0 • vae=mse