diff --git a/audio_diffusion/blocks.py b/audio_diffusion/blocks.py
index 6261a90..a4fbf05 100644
--- a/audio_diffusion/blocks.py
+++ b/audio_diffusion/blocks.py
@@ -52,7 +52,10 @@ def __init__(self, *main):
         self.main = nn.Sequential(*main)
 
     def forward(self, input):
-        return torch.cat([self.main(input), input], dim=1)
+        main = self.main(input)
+        pad = (0, (input.size(-1) - main.size(-1)))
+        main = torch.nn.functional.pad(main, pad)
+        return torch.cat([main, input], dim=1)
 
 class FourierFeatures(nn.Module):
     def __init__(self, in_features, out_features, std=1.):