diff --git a/modules/deeplearning/espcn.py b/modules/deeplearning/espcn.py
index b99843c65a4d280d31c85f81ae24ef8163ce6ce0..40a79ee04f43d00fe6ddbb50ab0681352ec0f9b0 100644
--- a/modules/deeplearning/espcn.py
+++ b/modules/deeplearning/espcn.py
@@ -172,7 +172,7 @@ class ESPCN:
 
         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=(32, 32, self.n_chans))
+        self.X_img = tf.keras.Input(shape=(32, 32, self.n_chans))
 
         self.inputs.append(self.X_img)
 
@@ -196,6 +196,7 @@ class ESPCN:
         # label = label[:, label_idx, :, :]
         label = label[:, label_idx, 3:67, 3:67]
         label = np.expand_dims(label, axis=3)
+        label = tf.image.resize(label, (32, 32))
 
         # data = data[:, data_idx, :, :]
         data = data[:, data_idx, 3:67, 3:67]
@@ -366,8 +367,8 @@ class ESPCN:
         # conv = tf.keras.layers.Conv2D((factor ** 2), 3, padding='same')(conv)
         print(conv.shape)
 
-        conv = tf.nn.depth_to_space(conv, factor)
-        print(conv.shape)
+        # conv = tf.nn.depth_to_space(conv, factor)
+        # print(conv.shape)
 
         self.logits = tf.keras.layers.Conv2D(1, kernel_size=3, strides=1, padding=padding, name='regression')(conv)