diff --git a/models/base.py b/models/base.py
index fee7bc7..82ade7e 100644
--- a/models/base.py
+++ b/models/base.py
@@ -72,11 +72,12 @@ class AEEncoder(tf.keras.Model):
 
 class Decoder(tf.keras.Model):
     
-    def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False):
+    def __init__(self, input_size, depths, dims, kernel_size=7, is_label=False, upscale=False):
         super(Decoder, self).__init__()
 
         self.dense = tf.keras.layers.Dense(tf.reduce_prod(input_size))
         self.reshape = tf.keras.layers.Reshape(input_size)
+        self.upscale = upscale
         
         if not is_label:
             self.decoder = ConvNeXtReverse(depths=list(reversed(depths)), dims=list(reversed(dims)), kernel_size=kernel_size)
@@ -94,6 +95,10 @@ class Decoder(tf.keras.Model):
         x = self.reshape(x)
         x = self.decoder(x, training=training)
 
+        if self.upscale:
+            x = tf.image.resize(x, size=(512, 256), method=tf.image.ResizeMethod.LANCZOS5)
+
+
         if len(inputs) == 2:
             if len(mask.shape) == 2:
                 x = x * mask[tf.newaxis,:,:,tf.newaxis]
diff --git a/models/cond_encoder.py b/models/cond_encoder.py
index 10d1b5f..6f5c391 100644
--- a/models/cond_encoder.py
+++ b/models/cond_encoder.py
@@ -8,7 +8,7 @@ class AutoEncoder(tf.keras.Model):
         super(AutoEncoder, self).__init__()
 
         self.encoder : AEEncoder = AEEncoder(config['latent_dim'], config['encoder_depths'], config['encoder_dims'], config['bottleneck_size'], config['kernel_size'])
-        self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'])
+        self.decoder : Decoder = Decoder(config['decoder_input'], config['decoder_depths'], config['decoder_dims'], kernel_size=config['kernel_size'], is_label=config['is_label'], upscale=config['upscale_output'])
 
 
     def call(self, inputs, training=None):