Change timestep device to cpu for xla (#11501)
* Change timestep device to cpu for xla
* Add all pipelines
* ruff format
* Apply style fixes
---------
Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
