From 26bb7fa0cb584554cf6f642dc53cd00cd0d56b3e Mon Sep 17 00:00:00 2001
From: Juan Acevedo <juancevedo@gmail.com>
Date: Mon, 13 Apr 2026 20:51:26 -0700
Subject: [PATCH] [ptxla] fix pytorch xla inference on TPUs. (#13463)

Co-authored-by: Juan Acevedo <jfacevedo@google.com>
---
 src/diffusers/pipelines/flux/pipeline_flux.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/diffusers/pipelines/flux/pipeline_flux.py b/src/diffusers/pipelines/flux/pipeline_flux.py
index be2bbe2acc..e125924adf 100644
--- a/src/diffusers/pipelines/flux/pipeline_flux.py
+++ b/src/diffusers/pipelines/flux/pipeline_flux.py
@@ -877,10 +877,7 @@ class FluxPipeline(
             self.scheduler.config.get("max_shift", 1.15),
         )
 
-        if XLA_AVAILABLE:
-            timestep_device = "cpu"
-        else:
-            timestep_device = device
+        timestep_device = device
         timesteps, num_inference_steps = retrieve_timesteps(
             self.scheduler,
             num_inference_steps,