[Flax MLM] Refactor run mlm with optax (#11745)

* refactor

* update

* update

* update

* refactor run mlm

* finalize

* refactor more

* fix typo

* update

* finish refactor

* modify run mlm

* Apply suggestions from code review

* Apply suggestions from code review

* Apply suggestions from code review

* small fixes

* upload

* upload

* finish run mlm script

Co-authored-by: Patrick von Platen <patrick@huggingface.co>
This commit is contained in:
Patrick von Platen
2021-05-19 12:00:58 +01:00
committed by GitHub
parent 43891be19b
commit 00440e350f
4 changed files with 277 additions and 210 deletions

View File

@@ -1,5 +1,5 @@
datasets >= 1.1.3
jax>=0.2.8
jaxlib>=0.1.59
git+https://github.com/google/flax.git
flax>=0.3.4
git+https://github.com/deepmind/optax.git