Refactor transforms.Fade on GPU computation (#1871)

torchaudio/transforms.py

@@ -1025,12 +1025,15 @@ class Fade(torch.nn.Module):
         """
         waveform_length = waveform.size()[-1]
         device = waveform.device
-        return self._fade_in(waveform_length).to(device) * \
-            self._fade_out(waveform_length).to(device) * waveform
+        return (
+            self._fade_in(waveform_length, device)
+            * self._fade_out(waveform_length, device)
+            * waveform
+        )
 
-    def _fade_in(self, waveform_length: int) -> Tensor:
-        fade = torch.linspace(0, 1, self.fade_in_len)
-        ones = torch.ones(waveform_length - self.fade_in_len)
+    def _fade_in(self, waveform_length: int, device: torch.device) -> Tensor:
+        fade = torch.linspace(0, 1, self.fade_in_len, device=device)
+        ones = torch.ones(waveform_length - self.fade_in_len, device=device)
 
         if self.fade_shape == "linear":
             fade = fade
@@ -1049,9 +1052,9 @@ class Fade(torch.nn.Module):
 
         return torch.cat((fade, ones)).clamp_(0, 1)
 
-    def _fade_out(self, waveform_length: int) -> Tensor:
-        fade = torch.linspace(0, 1, self.fade_out_len)
-        ones = torch.ones(waveform_length - self.fade_out_len)
+    def _fade_out(self, waveform_length: int, device: torch.device) -> Tensor:
+        fade = torch.linspace(0, 1, self.fade_out_len, device=device)
+        ones = torch.ones(waveform_length - self.fade_out_len, device=device)
 
         if self.fade_shape == "linear":
             fade = - fade + 1