atong01 / conditional-flow-matching

TorchCFM: a Conditional Flow Matching library
https://arxiv.org/abs/2302.00482
MIT License
1.23k stars 100 forks source link

Fix DDP Example for CIFAR10 by Using Epochs Only #147

Open ImahnShekhzadeh opened 1 hour ago

ImahnShekhzadeh commented 1 hour ago

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

atong01 commented 1 hour ago

Yes I agree with this change in principle. Happy to take a PR in this direction. Thanks for all your work on this example!