cast timesteps to int in flax scheduler too

kashif · kashif · commit 2cb7a3373e9b · 2022-09-26T11:25:26.000+02:00
diff --git a/src/diffusers/schedulers/scheduling_lms_discrete_flax.py b/src/diffusers/schedulers/scheduling_lms_discrete_flax.py
@@ -134,7 +134,7 @@ def set_timesteps(self, state: LMSDiscreteSchedulerState, num_inference_steps: i
 
         return state.replace(
             num_inference_steps=num_inference_steps,
-            timesteps=timesteps,
+            timesteps=timesteps.astype(int),
             derivatives=jnp.array([]),
             sigmas=sigmas,
         )

Original file line number	Diff line number	Diff line change
`@@ -134,7 +134,7 @@ def set_timesteps(self, state: LMSDiscreteSchedulerState, num_inference_steps: i`
`134`	`134`
`135`	`135`	`return state.replace(`
`136`	`136`	`num_inference_steps=num_inference_steps,`
`137`		`- timesteps=timesteps,`
	`137`	`+ timesteps=timesteps.astype(int),`
`138`	`138`	`derivatives=jnp.array([]),`
`139`	`139`	`sigmas=sigmas,`
`140`	`140`	`)`