Move WithResample layer to torch_wae.network.

src/torch_wae/cli/export_to_onnx.py

@@ -6,7 +6,7 @@
 import typer
 from torchaudio import functional as F
 
-from torch_wae.network import WAENet
+from torch_wae.network import WAENet, WithResample
 
 app = typer.Typer()
 
@@ -51,18 +51,5 @@ def main(
     )
 
 
-class WithResample(torch.nn.Module):
-    def __init__(self, f: WAENet, sample_rate: int) -> None:
-        super().__init__()
-
-        self.f = f
-        self.sample_rate = sample_rate
-
-    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
-        h = F.resample(waveform, self.sample_rate, self.f.SAMPLE_RATE)
-        z = self.f(h)
-        return z
-
-
 if __name__ == "__main__":
     app()

src/torch_wae/network.py

@@ -137,3 +137,16 @@ def __init__(self, dim: int = 1, eps: float = 1e-12):
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return F.normalize(x, p=2, dim=self.dim, eps=self.eps)
+
+
+class WithResample(torch.nn.Module):
+    def __init__(self, f: WAENet, sample_rate: int) -> None:
+        super().__init__()
+
+        self.f = f
+        self.sample_rate = sample_rate
+
+    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+        h = F.resample(waveform, self.sample_rate, self.f.SAMPLE_RATE)
+        z = self.f(h)
+        return z