Add missing set_timesteps in tests

Author: skirsten

tests/test_scheduler_flax.py

@@ -878,6 +878,9 @@ def test_inference_plms_no_past_residuals(self):
             scheduler = scheduler_class(**scheduler_config)
             state = scheduler.create_state()
 
+            num_inference_steps = 10
+            state = scheduler.set_timesteps(state, num_inference_steps, shape=self.dummy_sample.shape)
+
             scheduler.step_plms(state, self.dummy_sample, 1, self.dummy_sample).prev_sample
 
     def test_full_loop_no_noise(self):