diff --git a/modules/deeplearning/srcnn_cld_frac.py b/modules/deeplearning/srcnn_cld_frac.py
index 83fea1498c8484046015ddf4490b467508fcc2a4..fd24a93e90abc9506ec44d32729bab1c91def8e8 100644
--- a/modules/deeplearning/srcnn_cld_frac.py
+++ b/modules/deeplearning/srcnn_cld_frac.py
@@ -762,8 +762,9 @@ class SRCNN:
         preds = np.concatenate(self.test_preds)
         print(labels.shape, preds.shape)
 
-        labels_denorm = denormalize(labels, label_param, mean_std_dct)
-        preds_denorm = denormalize(preds, label_param, mean_std_dct)
+        if label_param != 'cloud_probability':
+            labels_denorm = denormalize(labels, label_param, mean_std_dct)
+            preds_denorm = denormalize(preds, label_param, mean_std_dct)
 
         return labels_denorm, preds_denorm
 
@@ -815,7 +816,7 @@ def run_restore_static(directory, ckpt_dir, out_file=None):
     nn = SRCNN()
     labels_denorm, preds_denorm = nn.run_restore(directory, ckpt_dir)
     if out_file is not None:
-        np.save(out_file, [labels_denorm, preds_denorm])
+        np.save(out_file, [np.squeeze(labels_denorm), preds_denorm.argmax(axis=3)])
 
 
 def run_evaluate_static(in_file, out_file, ckpt_dir):