diff --git a/modules/deeplearning/cloud_fraction_fcn_abi_hkm_refl.py b/modules/deeplearning/cloud_fraction_fcn_abi_hkm_refl.py
index 6d56d83069027b14b8c8cdce419d6f8b042ccb85..d4b2305384927af812c6c352dfe94e688ca5ec39 100644
--- a/modules/deeplearning/cloud_fraction_fcn_abi_hkm_refl.py
+++ b/modules/deeplearning/cloud_fraction_fcn_abi_hkm_refl.py
@@ -611,12 +611,13 @@ class SRCNN:
             self.train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='train_accuracy')
             self.test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name='test_accuracy')
 
-    @tf.function(input_signature=[tf.TensorSpec(None, tf.float32), tf.TensorSpec(None, tf.float32)])
+    # @tf.function(input_signature=[tf.TensorSpec(None, tf.float32), tf.TensorSpec(None, tf.float32)])
+    @tf.function
     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({'2km': inputs[0], 'hkm': inputs[1]}, training=True)
+            # pred = self.model(inputs, training=True)
+            pred = self.model({'input_1': inputs[0], 'input_2': inputs[1]}, training=True)
             loss = self.loss(labels, pred)
             total_loss = loss
             if len(self.model.losses) > 0:
@@ -632,11 +633,12 @@ class SRCNN:
 
         return loss
 
-    @tf.function(input_signature=[tf.TensorSpec(None, tf.float32), tf.TensorSpec(None, tf.float32)])
+    # @tf.function(input_signature=[tf.TensorSpec(None, tf.float32), tf.TensorSpec(None, tf.float32)])
+    @tf.function
     def test_step(self, inputs, labels):
         labels = tf.squeeze(labels, axis=[3])
-        # pred = self.model([inputs], training=False)
-        pred = self.model({'2km': inputs[0], 'hkm': inputs[1]}, training=False)
+        # pred = self.model(inputs, training=False)
+        pred = self.model({'input_1': inputs[0], 'input_2': inputs[1]}, training=False)
         t_loss = self.loss(labels, pred)
 
         self.test_loss(t_loss)
@@ -645,8 +647,8 @@ class SRCNN:
     # @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)
-        pred = self.model({'2km': inputs[0], 'hkm': inputs[1]}, training=False)
+        # pred = self.model(inputs, training=False)
+        pred = self.model({'input_1': inputs[0], 'input_2': inputs[1]}, training=False)
         # t_loss = self.loss(tf.squeeze(labels, axis=[3]), pred)
         t_loss = self.loss(labels, pred)