Merge pull request #26 from will-rice/wr-sane-lengths

Add support for sane sizes 16000, 32000, etc.
will-rice · May 11, 2024 · 308dea1 · 308dea1
2 parents 6904804 + b8c9eef
commit 308dea1
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/denoisers/modeling/unet1d/model.py b/denoisers/modeling/unet1d/model.py
@@ -245,7 +245,7 @@ def forward(self, inputs: Tensor) -> Tensor:
         out = self.middle(out)
 
         for skip, layer in zip(reversed(skips), self.decoder_layers):
-            out = layer(out + skip)
+            out = layer(out[..., : skip.size(-1)] + skip)
 
         out = torch.concat([out, inputs], dim=1)
         out = self.out_conv(out)

diff --git a/denoisers/modeling/waveunet/model.py b/denoisers/modeling/waveunet/model.py
@@ -247,7 +247,7 @@ def forward(self, inputs: Tensor) -> Tensor:
         out = self.middle(out)
 
         for skip, layer in zip(reversed(skips), self.decoder_layers):
-            out = torch.concat([out, skip], dim=1)
+            out = torch.concat([out[..., : skip.size(-1)], skip], dim=1)
             out = layer(out)
 
         out = torch.concat([out, inputs], dim=1)