# Set TPU device in Trainer

**URL:** <https://discuss.huggingface.co/t/set-tpu-device-in-trainer/5109>\
**Category:** Beginners\
**Created:** [March 26, 2021, 5:56am UTC](https://discuss.huggingface.co/t/set-tpu-device-in-trainer/5109 "2021-03-26T05:56:36Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![Constantin](https://avatars.discourse-cdn.com/v4/letter/c/74df32/32.png) [@Constantin](https://discuss.huggingface.co/u/Constantin)\
**Post date:** [March 26, 2021, 5:56am UTC](https://discuss.huggingface.co/t/set-tpu-device-in-trainer/5109/1 "2021-03-26T05:56:36Z")

</div>

Hi,

I want to use TPU provided by Kaggle in my project. I use PyTorch XLA to do that

```
import torch_xla
import torch_xla.core.xla_model as xm    
device = xm.xla_device()

```

Then I define a model

`model = AutoModelForMaskedLM.from_pretrained("xlm-roberta-base")`

And as I can see model on xla device, it fine

`model.device # device(type='xla', index=1)`

Then I define Trainer instance with my model

```
trainer = Trainer(
    model,
    args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=data_collator,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

```

And train it

`trainer.train()`

But it seems to me that trainer does not use xla device for TPU device is idle in Kaggle…

So, how to use TPU device in Kaggle with PyTorch XLA in my case?

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](https://discuss.huggingface.co/u/sgugger)\
**Post date:** [March 26, 2021, 12:53pm UTC](https://discuss.huggingface.co/t/set-tpu-device-in-trainer/5109/2 "2021-03-26T12:53:10Z")

</div>

No, the trainer does not support training with TPUS inside colab. You have to use it in a script (like the example scripts) and launch training with our launcher (see [here](https://github.com/huggingface/transformers/tree/master/examples#running-on-tpus) for the instructions).

---

<div class="post-metadata">

**Author:** ![Constantin](https://avatars.discourse-cdn.com/v4/letter/c/74df32/32.png) [@Constantin](https://discuss.huggingface.co/u/Constantin)\
**Post date:** [March 26, 2021, 7:48pm UTC](https://discuss.huggingface.co/t/set-tpu-device-in-trainer/5109/3 "2021-03-26T19:48:19Z")

</div>

Thank you very much!

Firstly, I have tried to do the example, but it does not work… I do it in Kaggle notebook (is it correct?) and the session stops because of memory use limit

 ![](https://us1.discourse-cdn.com/hellohellohello/original/2X/0/07830a8dfea2bcc9a9b0bc7ab305323be254680e.jpeg)

How much memory I am supposed to have? Is Kaggle okay for it?

---

<div class="post-metadata">

**Author:** ![sgugger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/sgugger/32/2291_2.png) [@sgugger](https://discuss.huggingface.co/u/sgugger)\
**Post date:** [March 26, 2021, 7:59pm UTC](https://discuss.huggingface.co/t/set-tpu-device-in-trainer/5109/4 "2021-03-26T19:59:49Z")

</div>

I don’t see which memory you have maxed out on your screenshot. Is it the TPU memory? That seems unlikely for this example and that batch size.

---

<div class="post-metadata">

**Author:** ![Constantin](https://avatars.discourse-cdn.com/v4/letter/c/74df32/32.png) [@Constantin](https://discuss.huggingface.co/u/Constantin)\
**Post date:** [March 27, 2021, 7:51am UTC](https://discuss.huggingface.co/t/set-tpu-device-in-trainer/5109/5 "2021-03-27T07:51:58Z")

</div>

Probably it is Kaggle thing question, but here is another screenshot. It looks to me that RAM is just can’t handle it.

 ![](https://us1.discourse-cdn.com/hellohellohello/original/2X/9/9fdf4bd6ce57102f1a75911d5433fd837b0cc88d.jpeg)

---

<div class="post-metadata">

**Author:** ![saeed899](https://avatars.discourse-cdn.com/v4/letter/s/3da27b/32.png) [@saeed899](https://discuss.huggingface.co/u/saeed899)\
**Post date:** [October 15, 2024, 6:17am UTC](https://discuss.huggingface.co/t/set-tpu-device-in-trainer/5109/6 "2024-10-15T06:17:33Z")

</div>

@sgugger  
recently Kaggle replaced TPU by TPU VMs.  
but still problem