Ability to pass custom optimizer

[EpicUsaMan]

There is few optimizers which can show better generalization and training time: AdEMAMix and AdEMAMixShampoo

It can be included in repo or add some interface to change optimizer by yourself