diff --git a/modules/deeplearning/unet.py b/modules/deeplearning/unet.py
index e16946ef7b6259fcc01c012c34f5f8f0e72ea22c..27fcc7bd457d8fb16f4d0e50d7ac7ff5bc6d0e67 100644
--- a/modules/deeplearning/unet.py
+++ b/modules/deeplearning/unet.py
@@ -610,22 +610,26 @@ class UNET:
         conv = tf.keras.layers.concatenate([conv, conv_4])
         conv = tf.keras.layers.Conv2D(num_filters, kernel_size=3, strides=1, padding=padding, activation=activation)(conv)
         conv = tf.keras.layers.BatchNormalization()(conv)
+        print(conv.shape)
 
         num_filters /= 2
         conv = tf.keras.layers.Conv2DTranspose(num_filters, kernel_size=3, strides=2, padding=padding)(conv)
         conv = tf.keras.layers.concatenate([conv, conv_3])
         conv = tf.keras.layers.Conv2D(num_filters, kernel_size=3, strides=1, padding=padding, activation=activation)(conv)
         conv = tf.keras.layers.BatchNormalization()(conv)
+        print(conv.shape)
 
         num_filters /= 2
         conv = tf.keras.layers.Conv2DTranspose(num_filters, kernel_size=3, strides=2, padding=padding)(conv)
         conv = tf.keras.layers.concatenate([conv, conv_2])
         conv = tf.keras.layers.Conv2D(num_filters, kernel_size=3, strides=1, padding=padding, activation=activation)(conv)
         conv = tf.keras.layers.BatchNormalization()(conv)
+        print(conv.shape)
 
         num_filters /= 2
         conv = tf.keras.layers.Conv2DTranspose(num_filters, kernel_size=3, strides=2, padding=padding)(conv)
         conv = tf.keras.layers.concatenate([conv, conv_1])
+        print(conv.shape)
 
         if NumClasses == 2:
             activation = tf.nn.sigmoid  # For binary