* Refactor optimize

* Add accelerate tests

* Remove line

* Make style

* Print out accelerate

* Fixup

* Fixup

* Fixup

* Add mixed precision with accelerate

* Fixup mixed precision

* Make style

* Print logs only for rank zero

* Allow tf32

* Fixup

* Fixup

* Fix setting up logger

* Update README

* Fix linter

* Fixup

* Update test

* Handle different types in prepare

* Compute grad norm

* Fix eval_epoch single run

* Make style

* Meh linter