diff --git a/modules/deeplearning/icing_fcn.py b/modules/deeplearning/icing_fcn.py
index cceb93adfe47054b2adb8eb031f9dac771332ab5..32d00b884878b653eba4fcee9a867d2edc069aa6 100644
--- a/modules/deeplearning/icing_fcn.py
+++ b/modules/deeplearning/icing_fcn.py
@@ -401,8 +401,8 @@ class IcingIntensityFCN:
         nda = np.zeros([nd_idxs.size])
         nda[:] = self.flight_level
         nda = tf.one_hot(nda, 5).numpy()
-        nda = np.expand_dims(nda, axis=0)
-        nda = np.expand_dims(nda, axis=0)
+        nda = np.expand_dims(nda, axis=1)
+        nda = np.expand_dims(nda, axis=1)
 
         return data, nda