@@ -99,7 +99,10 @@ def dummy_vqvae_and_unet(self):
     @slow
     def test_audio_diffusion(self):
         device = "cpu"  # ensure determinism for the device-dependent torch.Generator
-        mel = Mel()
+        mel = Mel(
+            x_res=self.dummy_unet.config.sample_size[1],
+            y_res=self.dummy_unet.config.sample_size[0],
+        )
 
         scheduler = DDPMScheduler()
         pipe = AudioDiffusionPipeline(vqvae=None, unet=self.dummy_unet, mel=mel, scheduler=scheduler)
@@ -127,6 +130,11 @@ def test_audio_diffusion(self):
         assert np.abs(image_slice.flatten() - expected_slice).max() == 0
         assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() == 0
 
+        mel = Mel(
+            x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1],
+            y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0],
+        )
+
         scheduler = DDIMScheduler()
         dummy_vqvae_and_unet = self.dummy_vqvae_and_unet
         pipe = AudioDiffusionPipeline(
@@ -154,13 +162,15 @@ def test_audio_diffusion(self):
         pipe = AudioDiffusionPipeline(
             vqvae=self.dummy_vqvae_and_unet[0], unet=dummy_unet_condition, mel=mel, scheduler=scheduler
         )
+        pipe = pipe.to(device)
+        pipe.set_progress_bar_config(disable=None)
 
         np.random.seed(0)
         encoding = torch.rand((1, 1, 10))
         output = pipe(generator=generator, encoding=encoding)
         image = output.images[0]
         image_slice = np.frombuffer(image.tobytes(), dtype="uint8")[:10]
-        expected_slice = np.array([120, 139, 147, 123, 124, 96, 115, 121, 126, 144])
+        expected_slice = np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111])
 
         assert np.abs(image_slice.flatten() - expected_slice).max() == 0