diff --git a/modules/deeplearning/cnn_cld_frac.py b/modules/deeplearning/cnn_cld_frac.py
index 31ad79bb887c87653f5422c32e8646f16fe36bd1..c023a17c2fc59a70db7b5c437e16df764f455630 100644
--- a/modules/deeplearning/cnn_cld_frac.py
+++ b/modules/deeplearning/cnn_cld_frac.py
@@ -523,9 +523,9 @@ class SRCNN:
 
         conv_b = build_residual_conv2d_block(conv_b, num_filters, 'Residual_Block_4', kernel_size=KERNEL_SIZE, scale=scale)
 
-        # conv_b = build_residual_conv2d_block(conv_b, num_filters, 'Residual_Block_5', kernel_size=KERNEL_SIZE, scale=scale)
+        conv_b = build_residual_conv2d_block(conv_b, num_filters, 'Residual_Block_5', kernel_size=KERNEL_SIZE, scale=scale)
 
-        # conv_b = build_residual_conv2d_block(conv_b, num_filters, 'Residual_Block_6', kernel_size=KERNEL_SIZE, scale=scale)
+        conv_b = build_residual_conv2d_block(conv_b, num_filters, 'Residual_Block_6', kernel_size=KERNEL_SIZE, scale=scale)
 
         conv_b = build_residual_block_conv2d_down2x(conv_b, num_filters, activation)
 
@@ -559,7 +559,7 @@ class SRCNN:
         # self.loss = tf.keras.losses.MeanAbsoluteError()  # Regression
 
         # decayed_learning_rate = learning_rate * decay_rate ^ (global_step / decay_steps)
-        initial_learning_rate = 0.006
+        initial_learning_rate = 0.002
         decay_rate = 0.95
         steps_per_epoch = int(self.num_data_samples/BATCH_SIZE)  # one epoch
         decay_steps = int(steps_per_epoch) * 4
@@ -595,12 +595,11 @@ class SRCNN:
             self.train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
             self.test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')
 
-    @tf.function
-    def train_step(self, mini_batch):
-        inputs = [mini_batch[0]]
-        labels = mini_batch[1]
+    @tf.function(input_signature=[tf.TensorSpec(None, tf.float32), tf.TensorSpec(None, tf.float32)])
+    def train_step(self, inputs, labels):
+        labels = tf.squeeze(labels, axis=[3])
         with tf.GradientTape() as tape:
-            pred = self.model(inputs, training=True)
+            pred = self.model([inputs], training=True)
             loss = self.loss(labels, pred)
             total_loss = loss
             if len(self.model.losses) > 0:
@@ -616,20 +615,20 @@ class SRCNN:
 
         return loss
 
-    @tf.function
-    def test_step(self, mini_batch):
-        inputs = [mini_batch[0]]
-        labels = mini_batch[1]
-        pred = self.model(inputs, training=False)
+    @tf.function(input_signature=[tf.TensorSpec(None, tf.float32), tf.TensorSpec(None, tf.float32)])
+    def test_step(self, inputs, labels):
+        labels = tf.squeeze(labels, axis=[3])
+        pred = self.model([inputs], training=False)
         t_loss = self.loss(labels, pred)
 
         self.test_loss(t_loss)
         self.test_accuracy(labels, pred)
 
-    def predict(self, mini_batch):
-        inputs = [mini_batch[0]]
-        labels = mini_batch[1]
-        pred = self.model(inputs, training=False)
+    # @tf.function(input_signature=[tf.TensorSpec(None, tf.float32), tf.TensorSpec(None, tf.float32)])
+    # decorator commented out because pred.numpy(): pred not evaluated yet.
+    def predict(self, inputs, labels):
+        pred = self.model([inputs], training=False)
+        # t_loss = self.loss(tf.squeeze(labels), pred)
         t_loss = self.loss(labels, pred)
 
         self.test_labels.append(labels)