apax-hub / apax

A flexible and performant framework for training machine learning potentials.
MIT License
15 stars 2 forks source link

EMA, SAM update, cyclic LR schedule #269

Closed M-R-Schaefer closed 5 months ago

M-R-Schaefer commented 5 months ago

Added EMA parameters, replaced my custom SAM implementation with the one in optax (faster and probably more correct) and added a cyclic cosine lr schedule.

todo