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

Co-authored-by: Juan Acevedo <jfacevedo@google.com>

Author: entrpn

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,