diff --git a/modules/deeplearning/cloud_opd_fcn_abi.py b/modules/deeplearning/cloud_opd_fcn_abi.py
index 7a0de0a1a74d7f6dac131632df3485a6d160f53d..df4c7d4e67da8a073145045dd0793b453bca3233 100644
--- a/modules/deeplearning/cloud_opd_fcn_abi.py
+++ b/modules/deeplearning/cloud_opd_fcn_abi.py
@@ -903,7 +903,7 @@ def run_restore_static(directory, ckpt_dir, out_file=None):
                  descale(inputs[:, 1:y_hi, 1:x_hi, 2], 'refl_0_65um_nom', mean_std_dct),
                  inputs[:, 1:y_hi, 1:x_hi, 3],
                  descale(inputs[:, 1:y_hi, 1:x_hi, 4], label_param, mean_std_dct),
-                 np.expand_dims(cat_cld_frac[:, 1:y_hi, 1:x_hi], axis=3)])
+                 cat_cld_frac[:, 1:y_hi, 1:x_hi]])
 
 
 def run_evaluate_static(in_file, out_file, ckpt_dir):