Fix DDP Example for CIFAR10 by Using Epochs Only

[ImahnShekhzadeh]

Hi!

As recently discussed in #145 and #144 with @Xiaoming-Zhao (and I as had already mentioned in https://github.com/atong01/conditional-flow-matching/pull/116#discussion_r1695722539), I/we believe that it would be best to use epochs instead of steps in the distributed mode, i.e. in the training loop of train_cifar10_ddp.py. Is there a reason not to (except consistency to the other scripts)? :upside_down_face:

I would be happy to make the changes, test them extensively and then open a PR.

Cheers, Imahn