Fix from_pretrained with corrupted state_dict (#12939)

* Fix from_pretrained with corrupted state_dict * Adapt test * Use better checkpoint * Style * Clean up
2021-08-04 11:48:39 +02:00
parent a28da4c490
commit d4c834d2e0
4 changed files with 9 additions and 5 deletions
--- a/tests/test_benchmark_tf.py
+++ b/tests/test_benchmark_tf.py
@@ -52,7 +52,7 @@ class TFBenchmarkTest(unittest.TestCase):
        self.check_results_dict_not_empty(results.memory_inference_result)

    def test_inference_no_configs_only_pretrain(self):
-        MODEL_ID = "sshleifer/tiny-distilbert-base-uncased-finetuned-sst-2-english"
+        MODEL_ID = "sgugger/tiny-distilbert-classification"
        benchmark_args = TensorFlowBenchmarkArguments(
            models=[MODEL_ID],
            training=False,