Thanks for an faster version of DiT implementation. Helped me a lot.

I noticed that the extract_features.py saves the features file with a name that is same across the GPUs when using DDP so only single GPU's output is saved. I just changed name of the saving file to {rank}-{train_steps} and solved the problem.