Add predict step accumulation (#7767)

* Add eval_accumulation_step and clean distributed eval

* Add TPU test

* Add TPU stuff

* Fix arg name

* Fix Seq2SeqTrainer

* Fix total_size

* Update src/transformers/trainer_pt_utils.py

Co-authored-by: Lysandre Debut <lysandre@huggingface.co>

* Doc and add test to TPU

* Add unit test

* Adapt name

Co-authored-by: Lysandre Debut <lysandre@huggingface.co>

examples/seq2seq/seq2seq_trainer.py

@@ -174,7 +174,7 @@ class Seq2SeqTrainer(Trainer):
             # Call forward again to get loss # TODO: avoidable?
             outputs = model(**inputs, use_cache=False)
             loss = self._compute_loss(outputs[1], labels_out)
-            loss = loss.mean().item()
+            loss = loss.mean().detach()
             if self.args.prediction_loss_only:
                 return (loss, None, None)
 

examples/test_xla_examples.py

@@ -81,3 +81,14 @@ class TorchXLAExamplesTests(unittest.TestCase):
 
             # Assert that the script takes less than 300 seconds to make sure it doesn't hang.
             self.assertLess(end - start, 300)
+
+    def test_trainer_tpu(self):
+        import xla_spawn
+
+        testargs = """
+            transformers/tests/test_trainer_tpu.py
+            --num_cores=8
+            transformers/tests/test_trainer_tpu.py
+            """.split()
+        with patch.object(sys, "argv", testargs):
+            xla_spawn.main()