diff --git a/modules/deeplearning/cnn_cld_frac.py b/modules/deeplearning/cnn_cld_frac.py
index a2f9f83104cae7d3c108a631f345b27d976ce0c3..2ad505f915d8413c045d3c00d5e54d4c020bc6bd 100644
--- a/modules/deeplearning/cnn_cld_frac.py
+++ b/modules/deeplearning/cnn_cld_frac.py
@@ -736,7 +736,7 @@ class SRCNN:
                 trn_ds = trn_ds.batch(BATCH_SIZE)
                 for mini_batch in trn_ds:
                     if self.learningRateSchedule is not None:
-                        loss = self.train_step(mini_batch)
+                        loss = self.train_step(mini_batch[0], mini_batch[1])
 
                     if (step % 100) == 0:
 
@@ -751,7 +751,7 @@ class SRCNN:
                             tst_ds = tf.data.Dataset.from_tensor_slices((data_tst, label_tst))
                             tst_ds = tst_ds.batch(BATCH_SIZE)
                             for mini_batch_test in tst_ds:
-                                self.test_step(mini_batch_test)
+                                self.test_step(mini_batch_test[0], mini_batch_test[1])
 
                         with self.writer_valid.as_default():
                             tf.summary.scalar('loss_val', self.test_loss.result(), step=step)
@@ -780,7 +780,7 @@ class SRCNN:
                 ds = tf.data.Dataset.from_tensor_slices((data, label))
                 ds = ds.batch(BATCH_SIZE)
                 for mini_batch in ds:
-                    self.test_step(mini_batch)
+                    self.test_step(mini_batch[0], mini_batch[1])
 
             print('loss, acc: ', self.test_loss.result().numpy(), self.test_accuracy.result().numpy())
             print('------------------------------------------------------')