diff --git a/modules/deeplearning/cloud_fraction_fcn_viirs.py b/modules/deeplearning/cloud_fraction_fcn_viirs.py
index cd5bf4b1aabf56d94fc4e7d8bd0824ea41e91780..29722007a88d0516db84b7142605a607a652a3a8 100644
--- a/modules/deeplearning/cloud_fraction_fcn_viirs.py
+++ b/modules/deeplearning/cloud_fraction_fcn_viirs.py
@@ -572,7 +572,7 @@ class SRCNN:
 
         self.test_labels.append(labels)
         self.test_preds.append(pred.numpy())
-        self.test_input.append(inputs[:, :, :, 4])
+        self.test_input.append(inputs)
 
         self.test_loss(t_loss)
         self.test_accuracy(labels, pred)
@@ -772,7 +772,7 @@ def run_restore_static(directory, ckpt_dir, out_file=None):
     labels, preds, inputs = nn.run_restore(directory, ckpt_dir)
     if out_file is not None:
         np.save(out_file,
-                [np.squeeze(labels), preds.argmax(axis=3), inputs])
+                [np.squeeze(labels), preds.argmax(axis=3), inputs[:, 1:65, 1:65, 4]])
 
 
 def run_evaluate_static(in_file, out_file, ckpt_dir):