Repository: pytorch/audio

Issue:
dither with noise_shaping=True not working on cuda input 
## 🐛 Bug

When using the functional `dither` with noise shaping, line https://github.com/pytorch/audio/blob/e808225f1859fbf47e778b73c64e82669d1c1e94/torchaudio/functional.py#L1660 doesn't get the device from the input waveform.

## To Reproduce

Steps to reproduce the behavior:
```
In [5]: F.dither(torch.randn(2, 441000).cuda())                                                                       
Out[5]: 
tensor([[ 0.3492, -1.8493, -0.4794,  ...,  0.3907, -0.4662, -0.2815],
        [-1.3821,  0.0401,  0.0023,  ..., -1.1417,  0.6411, -0.6851]],
       device='cuda:0')

In [6]: F.dither(torch.randn(2, 441000).cuda(), noise_shaping=True)                                                   
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-6-5dfd051b2e21> in <module>
----> 1 F.dither(torch.randn(2, 441000).cuda(), noise_shaping=True)

~/anaconda3/lib/python3.7/site-packages/torchaudio-0.7.0a0+8fdb8df-py3.7-linux-x86_64.egg/torchaudio/functional.py in dither(waveform, density_function, noise_shaping)
   1836 
   1837     if noise_shaping:
-> 1838         return _add_noise_shaping(dithered, waveform)
   1839     else:
   1840         return dithered

~/anaconda3/lib/python3.7/site-packages/torchaudio-0.7.0a0+8fdb8df-py3.7-linux-x86_64.egg/torchaudio/functional.py in _add_noise_shaping(dithered_waveform, waveform)
   1732     for index in range(error.size()[0]):
   1733         err = error[index]
-> 1734         error_offset = torch.cat((torch.zeros(1), err))
   1735         error[index] = error_offset[:waveform.size()[1]]
   1736 
```

## Expected behavior

Dither should work with both `cpu` and `cuda` input.

## Environment

 - What commands did you used to install torchaudio (conda/pip/build from source)? Source
 - If you are building from source, which commit is it? 8fdb8df
 - What does `torchaudio.__version__` print? '0.7.0a0+8fdb8df'

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

PyTorch version: 1.6.0
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Ubuntu 18.04.4 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: version 3.10.2

Python version: 3.7
Is CUDA available: Yes
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: GeForce GTX 1080 Ti
Nvidia driver version: 440.44
cuDNN version: Could not collect

Versions of relevant libraries:
[pip3] numpy==1.17.2
[pip3] pytorch-lightning==0.5.3.2
[pip3] torch==1.3.0
[pip3] torchaudio==0.2
[pip3] torchaudio-contrib==0.1
[pip3] torchparse==0.1
[pip3] torchvision==0.4.0
[conda] Could not collect




Discussion:

Resolving patch:
diff --git a/torchaudio/functional.py b/torchaudio/functional.py
--- a/torchaudio/functional.py
+++ b/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


Test patch:
diff --git a/test/torchaudio_unittest/torchscript_consistency_impl.py b/test/torchaudio_unittest/torchscript_consistency_impl.py
--- a/test/torchaudio_unittest/torchscript_consistency_impl.py
+++ b/test/torchaudio_unittest/torchscript_consistency_impl.py
@@ -194,6 +194,13 @@ def func(tensor):
         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")