diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
index 5edd951..8d0beef 100644
--- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -302,7 +302,7 @@ class Unet(nn.Module):
         self.out_dim = default(out_dim, default_out_dim)
 
         self.final_conv = nn.Sequential(
-            block_klass(dim, dim),
+            block_klass(dim * 2, dim),
             nn.Conv2d(dim, self.out_dim, 1)
         )
 
@@ -324,12 +324,13 @@ class Unet(nn.Module):
         x = self.mid_block2(x, t)
 
         for block1, block2, attn, upsample in self.ups:
-            x = torch.cat((x, h.pop()), dim=1)
+            x = torch.cat((x, h.pop()), dim = 1)
             x = block1(x, t)
             x = block2(x, t)
             x = attn(x)
             x = upsample(x)
 
+        x = torch.cat((x, h.pop()), dim = 1)
         return self.final_conv(x)
 
 # gaussian diffusion trainer class
diff --git a/setup.py b/setup.py
index 3436bd6..3e4a964 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.18.4',
+  version = '0.19.0',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',