Fix dither with noise_shaping = True (#865)

test/torchaudio_unittest/torchscript_consistency_impl.py

@@ -194,6 +194,13 @@ class Functional(common_utils.TestBaseMixin):
         tensor = common_utils.get_whitenoise(n_channels=2)
         self._assert_consistency(func, tensor, shape_only=True)
 
+    def test_dither_noise_shaping(self):
+        def func(tensor):
+            return F.dither(tensor, noise_shaping=True)
+
+        tensor = common_utils.get_whitenoise(n_channels=2)
+        self._assert_consistency(func, tensor)
+
     def test_lfilter(self):
         if self.dtype == torch.float64:
             raise unittest.SkipTest("This test is known to fail for float64")

torchaudio/functional.py

@@ -1655,9 +1655,10 @@ def _add_noise_shaping(
     error = dithered_waveform - waveform
 
     # add error[n-1] to dithered_waveform[n], so offset the error by 1 index
+    zeros = torch.zeros(1, dtype=error.dtype, device=error.device)
     for index in range(error.size()[0]):
         err = error[index]
-        error_offset = torch.cat((torch.zeros(1), err))
+        error_offset = torch.cat((zeros, err))
         error[index] = error_offset[:waveform.size()[1]]
 
     noise_shaped = dithered_waveform + error