diff --git a/models/pyramid_level.py b/models/pyramid_level.py
index bc76f65..5d31bfe 100644
--- a/models/pyramid_level.py
+++ b/models/pyramid_level.py
@@ -9,7 +9,7 @@ import scipy as sp
 
 from custom_utils.custom_utils import load_model
 from models.base import AEEncoder, VariationalEncoder, Decoder
-from models.upscaler import UpscalerBasic, UpscalerBasicNew
+from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
 
 
 class PyramidLevel(tf.keras.Model):
@@ -24,7 +24,7 @@ class PyramidLevel(tf.keras.Model):
 
         self.upscale_enhancer = None
         if config['upscale_enhancer'] != '':
-            self.upscale_enhancer = UpscalerBasicNew(config['resolution'])
+            self.upscale_enhancer = UpscalerBasicConv(config['resolution'])
             load_model(self.upscale_enhancer, config['upscale_enhancer'])
             self.upscale_enhancer.trainable = not config['upscale_enhancer_freeze']
 
diff --git a/models/pyramid_model.py b/models/pyramid_model.py
index 0885f03..6f88345 100644
--- a/models/pyramid_model.py
+++ b/models/pyramid_model.py
@@ -10,7 +10,7 @@ from custom_utils.custom_utils import load_model
 from models.base import AEEncoder
 from models.cvae_model import CVAE
 from models.pyramid_level import PyramidLevel
-from models.upscaler import UpscalerBasic, UpscalerBasicNew
+from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
 
 
 class Pyramid(tf.keras.Model):
@@ -22,7 +22,7 @@ class Pyramid(tf.keras.Model):
         if config['pyramid']['static_model'] != '':
             load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
 
-        self.body_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
+        self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
         if config['pyramid']['body_model'] != '':
             load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
 
@@ -45,7 +45,7 @@ class Pyramid(tf.keras.Model):
  
         self.upscale_enhancer = None
         if config['pyramid']['final_upscale_enhancer'] != '':
-            self.upscale_enhancer = UpscalerBasicNew(low_res=config['pyramid']['resolutions'][-1])
+            self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1])
             load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
             self.upscale_enhancer.trainable = False
 
diff --git a/models/upscaler.py b/models/upscaler.py
index 100ebf9..888007b 100644
--- a/models/upscaler.py
+++ b/models/upscaler.py
@@ -125,19 +125,19 @@ class UpscalerBasicNew(tf.keras.Model):
 
 class UpscalerBasicConv(tf.keras.Model):
     
-    def __init__(self, low_res=32):
+    def __init__(self, low_res=32, kernel_size=3):
         super(UpscalerBasicConv, self).__init__()
 
         self.low_res = low_res
 
         self.seq = tf.keras.Sequential([
-                    tf.keras.layers.Conv2D(4096, kernel_size=3, padding='same', activation='relu', name='upscale_enhancer_conv_1'),
+                    tf.keras.layers.Conv2D(256, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_1'),
                     tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_1'),
-                    tf.keras.layers.Conv2D(2048, kernel_size=3, padding='same', activation='relu', name='upscale_enhancer_conv_2'),
+                    tf.keras.layers.Conv2D(128, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_2'),
                     tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_2'),
-                    tf.keras.layers.Conv2D(1024, kernel_size=3, padding='same', activation='relu', name='upscale_enhancer_conv_3'),
+                    tf.keras.layers.Conv2D(64, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_3'),
                     tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_3'),
-                    tf.keras.layers.Conv2D(128, kernel_size=3, padding='same', activation='relu', name='upscale_enhancer_conv_4'),
+                    tf.keras.layers.Conv2D(32, kernel_size=kernel_size, padding='same', activation='relu', name='upscale_enhancer_conv_4'),
                     tf.keras.layers.BatchNormalization(name='upscale_enhancer_batch_norm_4'),
                     tf.keras.layers.Dense(3, activation='tanh', name='upscale_enhancer_dense_final'),
         ])