test_add_noise_device

huggingface · Jan 8, 2025 · ac2b820 · ac2b820
1 parent f7fb73e
commit ac2b820
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/schedulers/test_schedulers.py b/tests/schedulers/test_schedulers.py
@@ -723,9 +723,9 @@ def test_add_noise_device(self):
             self.assertEqual(sample.shape, scaled_sample.shape)
 
             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)
+            t = scheduler.timesteps[5].expand(noise.shape[0])
+            noised = scheduler.add_noise(scaled_sample, noise, t)
+            self.assertEqual(noised.shape, scaled_sample.shape)
 
     def test_deprecated_kwargs(self):
         for scheduler_class in self.scheduler_classes: