diff --git a/modules/deeplearning/unet.py b/modules/deeplearning/unet.py
index 27fcc7bd457d8fb16f4d0e50d7ac7ff5bc6d0e67..d0c23939caa6b0e53387c8b76465e7e27aed4e6f 100644
--- a/modules/deeplearning/unet.py
+++ b/modules/deeplearning/unet.py
@@ -206,7 +206,8 @@ class UNET:
         n_chans = len(self.train_params)
         if TRIPLET:
             n_chans *= 3
-        self.X_img = tf.keras.Input(shape=(None, None, n_chans))
+        #self.X_img = tf.keras.Input(shape=(None, None, n_chans))
+        self.X_img = tf.keras.Input(shape=(16, 16, n_chans))
 
         self.inputs.append(self.X_img)
         self.inputs.append(tf.keras.Input(shape=(None, None, 5)))
@@ -604,6 +605,7 @@ class UNET:
         conv = tf.keras.layers.LeakyReLU()(conv)
 
         # Expanding (Decoding) branch -------------------------------------------------------------------------------
+        print('expanding branch')
 
         num_filters /= 2
         conv = tf.keras.layers.Conv2DTranspose(num_filters, kernel_size=3, strides=2, padding=padding)(conv)