diff --git a/audio_diffusion_pytorch/modules.py b/audio_diffusion_pytorch/modules.py
index 82213fb..4083dd5 100644
--- a/audio_diffusion_pytorch/modules.py
+++ b/audio_diffusion_pytorch/modules.py
@@ -1165,7 +1165,7 @@ def __init__(
         )
 
         self.to_out = Unpatcher(
-            in_channels=channels,
+            in_channels=channels * (use_noisy + 1),
             out_channels=in_channels,
             blocks=patch_blocks,
             factor=patch_factor,
@@ -1188,6 +1188,8 @@ def decode(self, x: Tensor) -> Tensor:
             if self.use_noisy:
                 x = torch.cat([x, torch.randn_like(x)], dim=1)
             x = upsample(x)
+        if self.use_noisy:
+            x = torch.cat([x, torch.randn_like(x)], dim=1)
         return self.to_out(x)
 
 
diff --git a/setup.py b/setup.py
index 6846dd6..8ba7eae 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@
 setup(
     name="audio-diffusion-pytorch",
     packages=find_packages(exclude=[]),
-    version="0.0.51",
+    version="0.0.52",
     license="MIT",
     description="Audio Diffusion - PyTorch",
     long_description_content_type="text/markdown",