# Pipeline device issue, torch\_xla generation() bug, flax models malloc errors **URL:** \ **Category:** 🤗Transformers\ **Created:** [April 21, 2024, 7:17pm UTC](https://discuss.huggingface.co/t/pipeline-device-issue-torch-xla-generation-bug-flax-models-malloc-errors/82961 "2024-04-21T19:17:27Z")\ **Posts on this page:** 1\ **Page:** 1
**Author:** ![vivek9840](https://avatars.discourse-cdn.com/v4/letter/v/eb8c5e/32.png) [@vivek9840](https://discuss.huggingface.co/u/vivek9840)\ **Post date:** [April 21, 2024, 7:17pm UTC](https://discuss.huggingface.co/t/pipeline-device-issue-torch-xla-generation-bug-flax-models-malloc-errors/82961/1 "2024-04-21T19:17:27Z")
I am inferring llms for certain text generation tasks with tpu. i originally wrote code using transformers pipeline with device\_map set to “auto”. but it never pick up the tpu device. then i give it as xm.xla\_device() as parameters but after doing that system, automatically crashes. after that i wrote code for generation. the models loaded on device but the generation was happening on cpu. this problem persist for other users too. referring to this [Generate text with `model.generate` on TPU does not work · Issue #12322 · huggingface/transformers · GitHub](https://github.com/huggingface/transformers/issues/12322) . then i tried to change torch model to flax. i used FlaxAutoModelForCausalLM for converting pytorch model to flax model. since flax version where not available. but while converting torch model to flax model. i got error for XLA buffer. which stated i had 29.7 MB space and i need 32.0 MB space for buffer. i tried to change environment variable for it but the buffer size did not changed. now i am currently trying to use jax models which where re-created by this person from pytoch to jax. [GitHub - ayaka14732/llama-2-jax: JAX implementation of the Llama 2 model](https://github.com/ayaka14732/llama-2-jax) . thank you.