fix bug in dpm_solver_jax.py

LuChengTHU / dpm-solver

Official code for "DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps" (Neurips 2022 Oral)

MIT License

1.52k stars 120 forks source link

Closed shaoshitong closed 1 year ago

shaoshitong commented 1 year ago

Due to jax version change, the code in dpm_solver_jax.py was not working, so in this PR I have made a code fix.

LuChengTHU commented 1 year ago

Thanks a lot!