test_add_noise_device

Author: hlky

tests/schedulers/test_schedulers.py

@@ -711,16 +711,16 @@ def test_add_noise_device(self):
             scheduler = scheduler_class(**scheduler_config)
             scheduler.set_timesteps(self.default_num_inference_steps)
 
-            # sample = self.dummy_sample.to(torch_device)
-            # if scheduler_class == CMStochasticIterativeScheduler:
-            #     # Get valid timestep based on sigma_max, which should always be in timestep schedule.
-            #     scaled_sigma_max = scheduler.sigma_to_t(scheduler.config.sigma_max)
-            #     scaled_sample = scheduler.scale_model_input(sample, scaled_sigma_max)
-            # elif scheduler_class == EDMEulerScheduler:
-            #     scaled_sample = scheduler.scale_model_input(sample, scheduler.timesteps[-1])
-            # else:
-            #     scaled_sample = scheduler.scale_model_input(sample, 0.0)
-            # self.assertEqual(sample.shape, scaled_sample.shape)
+            sample = self.dummy_sample.to(torch_device)
+            if scheduler_class == CMStochasticIterativeScheduler:
+                # Get valid timestep based on sigma_max, which should always be in timestep schedule.
+                scaled_sigma_max = scheduler.sigma_to_t(scheduler.config.sigma_max)
+                scaled_sample = scheduler.scale_model_input(sample, scaled_sigma_max)
+            elif scheduler_class == EDMEulerScheduler:
+                scaled_sample = scheduler.scale_model_input(sample, scheduler.timesteps[-1])
+            else:
+                scaled_sample = scheduler.scale_model_input(sample, 0.0)
+            self.assertEqual(sample.shape, scaled_sample.shape)
 
             # noise = torch.randn_like(scaled_sample).to(torch_device)
             # t = scheduler.timesteps[5][None]