[ptxla] fix pytorch xla inference on TPUs. (huggingface#13463)

entrpn · jfacevedo-google · web-flow · commit 26bb7fa0cb58 · 2026-04-14T09:21:26.000+05:30
Co-authored-by: Juan Acevedo &lt;jfacevedo@google.com&gt;
diff --git a/src/diffusers/pipelines/flux/pipeline_flux.py b/src/diffusers/pipelines/flux/pipeline_flux.py
@@ -877,10 +877,7 @@ def __call__(
             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,
