GitLab

feat(training-utils): support device and dtype params in compute_density_for_timestep_sampling (#10699)

* feat(training-utils): support device and dtype params in compute_density_for_timestep_sampling

* chore: update type hint

* refactor: use union for type hint

---------
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>