diff --git a/modules/deeplearning/espcn.py b/modules/deeplearning/espcn.py
index 36e529040d52357faf211b0a6ae588008481b38f..773d2108dd4729cb3de143a87be63e402ef7819a 100644
--- a/modules/deeplearning/espcn.py
+++ b/modules/deeplearning/espcn.py
@@ -211,7 +211,7 @@ class ESPCN:
         self.n_chans = 1
 
         self.X_img = tf.keras.Input(shape=(None, None, self.n_chans))
-        # self.X_img = tf.keras.Input(shape=(36, 36, self.n_chans))
+        self.X_img = tf.keras.Input(shape=(36, 36, self.n_chans))
         # self.X_img = tf.keras.Input(shape=(32, 32, self.n_chans))
 
         self.inputs.append(self.X_img)
@@ -235,12 +235,14 @@ class ESPCN:
         data = np.concatenate(label_s)
         label = np.concatenate(label_s)
 
-        label = label[:, label_idx, :, :]
+        # label = label[:, label_idx, :, :]
+        label = label[:, label_idx, 4:68, 4:68]
         label = np.expand_dims(label, axis=3)
 
         data = data[:, data_idx, :, :]
         data = np.expand_dims(data, axis=3)
-        data = tf.image.resize(data, (32, 32)).numpy()
+        # data = tf.image.resize(data, (32, 32)).numpy()
+        data = tf.image.resize(data, (36, 36)).numpy()
 
         data = data.astype(np.float32)
         label = label.astype(np.float32)
@@ -381,8 +383,8 @@ class ESPCN:
 
         input_2d = self.inputs[0]
         print('input: ', input_2d.shape)
-        # conv = tf.keras.layers.Conv2D(num_filters, kernel_size=5, strides=1, padding='VALID', activation=None)(input_2d)
-        conv = input_2d
+        conv = tf.keras.layers.Conv2D(num_filters, kernel_size=5, strides=1, padding='VALID', activation=None)(input_2d)
+        # conv = input_2d
         print('input: ', conv.shape)
         skip = conv
 
@@ -394,7 +396,7 @@ class ESPCN:
         if do_batch_norm:
             conv = tf.keras.layers.BatchNormalization()(conv)
 
-        conv = tf.keras.layers.Conv2D(num_filters, kernel_size=5, strides=1, padding=padding, activation=activation)(conv)
+        conv = tf.keras.layers.Conv2D(num_filters, kernel_size=3, strides=1, padding=padding, activation=activation)(conv)
         print(conv.shape)
 
         if do_drop_out:
@@ -439,6 +441,9 @@ class ESPCN:
         conv = tf.keras.layers.Conv2DTranspose(num_filters // 4, kernel_size=3, strides=1, padding=padding, activation=activation)(conv)
         print(conv.shape)
 
+        conv = tf.keras.layers.Conv2DTranspose(num_filters // 4, kernel_size=3, strides=1, padding=padding, activation=activation)(conv)
+        print(conv.shape)
+
         #self.logits = tf.keras.layers.Conv2D(1, kernel_size=1, strides=1, padding=padding, name='probability', activation=tf.nn.sigmoid)(conv)
         self.logits = tf.keras.layers.Conv2D(1, kernel_size=1, strides=1, padding=padding, name='probability')(conv)