[Flax examples] remove dependancy on pytorch training args (#14636)

* use custom training arguments * update tests
2021-12-12 09:19:12 +05:30
parent 027074f4d0
commit 6a025487a6
8 changed files with 457 additions and 17 deletions
--- a/examples/flax/test_examples.py
+++ b/examples/flax/test_examples.py
@@ -137,7 +137,7 @@ class ExamplesTests(TestCasePlus):
            --test_file tests/fixtures/tests_samples/xsum/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
-            --max_steps=50
+            --num_train_epochs=3
            --warmup_steps=8
            --do_train
            --do_eval
@@ -257,7 +257,7 @@ class ExamplesTests(TestCasePlus):
            --validation_file tests/fixtures/tests_samples/SQUAD/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
-            --max_steps=10
+            --num_train_epochs=3
            --warmup_steps=2
            --do_train
            --do_eval