❓ How to use TFTrainer on TPU ? Unable to destroy remote tensor handles
# ❓ Questions & Help

I'm trying to train my model on TPU using `TFTrainer`. Training starts fine, but after a few training steps, I'm having this error : 

> Unable to destroy remote tensor handles. If you are running a tf.function, it usually indicates some op in the graph gets an error: 4 root error(s) found.

I don't know what I am doing wrong, any help is greatly appreciated.

---

Here is the full stacktrace :

```
2020/06/15 05:14:31 - INFO - transformers.trainer_tf -   Epoch 2 Step 1500 Train Loss 3.0905
2020/06/15 05:17:33 - INFO - transformers.trainer_tf -   Epoch 2 Step 2000 Train Loss 2.8351
2020/06/15 05:20:35 - INFO - transformers.trainer_tf -   Epoch 2 Step 2500 Train Loss 3.1043
2020-06-15 05:22:53.919745: W tensorflow/core/distributed_runtime/eager/remote_tensor_handle_data.cc:76] Unable to destroy remote tensor handles. If you are running a tf.function, it usually indicates some op in the graph gets an error: 4 root error(s) found.
  (0) Cancelled: {{function_node __inference__accumulate_next_440382}} Function was cancelled before it was started
  (1) Cancelled: {{function_node __inference__accumulate_next_440382}} Function was cancelled before it was started
  (2) Cancelled: {{function_node __inference__accumulate_next_440382}} Function was cancelled before it was started
  (3) Out of range: {{function_node __inference__accumulate_next_440382}} End of sequence
         [[{{node IteratorGetNext_6}}]]
0 successful operations.
5 derived errors ignored.
Traceback (most recent call last):
  File "train.py", line 136, in <module>
    main()
  File "train.py", line 112, in main
    trainer.train()
  File "/home/me/.venv/x/lib/python3.6/site-packages/transformers/trainer_tf.py", line 274, in train
    for training_loss in self._training_steps(train_ds, optimizer):
  File "/home/me/.venv/x/lib/python3.6/site-packages/transformers/trainer_tf.py", line 319, in _training_steps
    self._apply_gradients(optimizer)
  File "/home/me/.venv/x/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 580, in __call__
    result = self._call(*args, **kwds)
  File "/home/me/.venv/x/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 611, in _call
    return self._stateless_fn(*args, **kwds)  # pylint: disable=not-callable
  File "/home/me/.venv/x/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 2420, in __call__
    return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
  File "/home/me/.venv/x/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1665, in _filtered_call
    self.captured_inputs)
  File "/home/me/.venv/x/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1746, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "/home/me/.venv/x/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 598, in call
    ctx=ctx)
  File "/home/me/.venv/x/lib/python3.6/site-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.OutOfRangeError: 4 root error(s) found.
  (0) Cancelled: {{function_node __inference__accumulate_next_440382}} Function was cancelled before it was started
  (1) Cancelled: {{function_node __inference__accumulate_next_440382}} Function was cancelled before it was started
  (2) Cancelled: {{function_node __inference__accumulate_next_440382}} Function was cancelled before it was started
  (3) Out of range: {{function_node __inference__accumulate_next_440382}} End of sequence
         [[{{node IteratorGetNext_6}}]]
0 successful operations.
5 derived errors ignored. [Op:__inference__apply_gradients_373775]
```
It seems to be related to the size of dataset : if I use `--training_steps` instead of `--num_epochs` I don't have the error.

---

But a similar error appear at evaluation time :

```
tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node __inference__evaluate_steps_517418}} Compilation failure: Output shapes of then and else branches do not match: (pred[1]) vs. (pred[])
         [[{{node lossed_bart/model/decoder/cond}}]]
        TPU compilation failed
         [[tpu_compile_succeeded_assert/_17738495357405513889/_9]]
```<|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
