diff --git a/denoising_diffusion_pytorch/classifier_free_guidance.py b/denoising_diffusion_pytorch/classifier_free_guidance.py
index 19bbbd3d4..d836a8762 100644
--- a/denoising_diffusion_pytorch/classifier_free_guidance.py
+++ b/denoising_diffusion_pytorch/classifier_free_guidance.py
@@ -349,8 +349,8 @@ def __init__(
         default_out_dim = channels * (1 if not learned_variance else 2)
         self.out_dim = default(out_dim, default_out_dim)
 
-        self.final_res_block = ResnetBlock(dim * 2, dim, time_emb_dim = time_dim, classes_emb_dim = classes_dim)
-        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
+        self.final_res_block = ResnetBlock(init_dim * 2, init_dim, time_emb_dim = time_dim, classes_emb_dim = classes_dim)
+        self.final_conv = nn.Conv2d(init_dim, self.out_dim, 1)
 
     def forward_with_cond_scale(
         self,
diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
index 4a69683ea..6a9f7da92 100644
--- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
+++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
@@ -380,8 +380,8 @@ def __init__(
         default_out_dim = channels * (1 if not learned_variance else 2)
         self.out_dim = default(out_dim, default_out_dim)
 
-        self.final_res_block = resnet_block(dim * 2, dim)
-        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
+        self.final_res_block = resnet_block(init_dim * 2, init_dim)
+        self.final_conv = nn.Conv2d(init_dim, self.out_dim, 1)
 
     @property
     def downsample_factor(self):
diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py
index eda42a105..e94dde819 100644
--- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py
+++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py
@@ -343,8 +343,8 @@ def __init__(
         default_out_dim = channels * (1 if not learned_variance else 2)
         self.out_dim = default(out_dim, default_out_dim)
 
-        self.final_res_block = resnet_block(dim * 2, dim)
-        self.final_conv = nn.Conv1d(dim, self.out_dim, 1)
+        self.final_res_block = resnet_block(init_dim * 2, init_dim)
+        self.final_conv = nn.Conv1d(init_dim, self.out_dim, 1)
 
     def forward(self, x, time, x_self_cond = None):
         if self.self_condition:
diff --git a/denoising_diffusion_pytorch/repaint.py b/denoising_diffusion_pytorch/repaint.py
index 395159984..475ac0930 100644
--- a/denoising_diffusion_pytorch/repaint.py
+++ b/denoising_diffusion_pytorch/repaint.py
@@ -376,8 +376,8 @@ def __init__(
         default_out_dim = channels * (1 if not learned_variance else 2)
         self.out_dim = default(out_dim, default_out_dim)
 
-        self.final_res_block = ResnetBlock(dim * 2, dim, time_emb_dim = time_dim)
-        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
+        self.final_res_block = ResnetBlock(init_dim * 2, init_dim, time_emb_dim = time_dim)
+        self.final_conv = nn.Conv2d(init_dim, self.out_dim, 1)
 
     @property
     def downsample_factor(self):
diff --git a/denoising_diffusion_pytorch/simple_diffusion.py b/denoising_diffusion_pytorch/simple_diffusion.py
index 26bbad00e..6be766ca7 100644
--- a/denoising_diffusion_pytorch/simple_diffusion.py
+++ b/denoising_diffusion_pytorch/simple_diffusion.py
@@ -428,8 +428,8 @@ def __init__(
         default_out_dim = input_channels
         self.out_dim = default(out_dim, default_out_dim)
 
-        self.final_res_block = ResnetBlock(dim * 2, dim, time_emb_dim = time_dim)
-        self.final_conv = nn.Conv2d(dim, self.out_dim, 1)
+        self.final_res_block = ResnetBlock(init_dim * 2, init_dim, time_emb_dim = time_dim)
+        self.final_conv = nn.Conv2d(init_dim, self.out_dim, 1)
 
     def forward(self, x, time):
         x = self.init_img_transform(x)
diff --git a/denoising_diffusion_pytorch/version.py b/denoising_diffusion_pytorch/version.py
index fca23f681..2c488f9d1 100644
--- a/denoising_diffusion_pytorch/version.py
+++ b/denoising_diffusion_pytorch/version.py
@@ -1 +1 @@
-__version__ = '2.0.10'
+__version__ = '2.0.12'