test_add_noise_device

huggingface · Jan 8, 2025 · f7fb73e · f7fb73e
1 parent 58b6081
commit f7fb73e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/schedulers/test_schedulers.py b/tests/schedulers/test_schedulers.py
@@ -722,8 +722,8 @@ def test_add_noise_device(self):
                 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]
+            noise = torch.randn_like(scaled_sample).to(torch_device)
+            t = scheduler.timesteps[5][None]
             # noised = scheduler.add_noise(scaled_sample, noise, t)
             # self.assertEqual(noised.shape, scaled_sample.shape)