❓ [TFTrainer] How to run on 8 TPU cores ?
# ❓ Questions & Help

I'm trying to run TFTrainer on 8 TPU cores, but I don't understand how to make it work.

I tried running my script with the flags `--tpu_num_cores 8 --per_device_train_batch_size 8`, expecting each core to handle a batch size of 1.

But when I print the shape of my inputs, I have `[8, x]` instead of `[1, x]`, which leads to memory error.

---

If I start the training with `--tpu_num_cores 8 --per_device_train_batch_size 1`, the shape of inputs is correct (`[1, x]`), but the number of optimization steps computed is not correct (If I have 8k samples, it says I have 8k optimization steps, but I expected 1k steps because I am using 8 TPU cores...).

---

Am I doing something wrong ? **How can I train on 8 TPU cores, with a batch size of 1 for each core ?**
Hello !

This is because the `--tpu_num_cores` is not taken into account yet. If you want to use TPUs, just fill the TPU name with `--tpu_name` and it will detect automatically the number of cores. For now TPUs with TF Trainer is under development so some use cases might not work properly.<|||||>Thanks for the input @jplu 

So as you mentioned the number of TPU cores is automatically detected. It is accessible with `training_args.n_gpu`.

Also I didn't notice but here : https://github.com/huggingface/transformers/blob/e4aaa4580515446cd5a2972ab42fec0b95819c84/src/transformers/training_args.py#L150

The batch size is automatically adjusted to the number of cores.

So the behavior I observed is completely normal, as `--per_device_train_batch_size` is the batch size **per TPU cores**. 

