diff --git a/modules/deeplearning/icing_fcn.py b/modules/deeplearning/icing_fcn.py
index 1803fc4698806b876ee0e7d32ec96b60e388cf29..576e3de0df6e96a54ed2c394fc57c37d534778f3 100644
--- a/modules/deeplearning/icing_fcn.py
+++ b/modules/deeplearning/icing_fcn.py
@@ -19,6 +19,7 @@ if NumClasses == 2:
     NumLogits = 1
 else:
     NumLogits = NumClasses
+NumFlightLevels = 3
 
 BATCH_SIZE = 128
 NUM_EPOCHS = 60
@@ -241,8 +242,7 @@ class IcingIntensityFCN:
         self.X_img = tf.keras.Input(shape=(None, None, n_chans))
 
         self.inputs.append(self.X_img)
-        #self.inputs.append(tf.keras.Input(shape=(None, None, 5)))
-        self.inputs.append(tf.keras.Input(shape=(None, None, 3)))
+        self.inputs.append(tf.keras.Input(shape=(None, None, NumFlightLevels)))
 
         self.flight_level = 0
 
@@ -362,18 +362,18 @@ class IcingIntensityFCN:
 
         nda = h5f[param][nd_idxs,]
 
-        # nda[np.logical_and(nda >= 0, nda < 2000)] = 0
-        # nda[np.logical_and(nda >= 2000, nda < 4000)] = 1
-        # nda[np.logical_and(nda >= 4000, nda < 6000)] = 2
-        # nda[np.logical_and(nda >= 6000, nda < 8000)] = 3
-        # nda[np.logical_and(nda >= 8000, nda < 15000)] = 4
-
-        nda[np.logical_and(nda >= 0, nda < 3000)] = 0
-        nda[np.logical_and(nda >= 3000, nda < 6000)] = 1
-        nda[np.logical_and(nda >= 6000, nda < 15000)] = 2
-
-        # nda = tf.one_hot(nda, 5).numpy()
-        nda = tf.one_hot(nda, 3).numpy()
+        if NumFlightLevels == 5:
+            nda[np.logical_and(nda >= 0, nda < 2000)] = 0
+            nda[np.logical_and(nda >= 2000, nda < 4000)] = 1
+            nda[np.logical_and(nda >= 4000, nda < 6000)] = 2
+            nda[np.logical_and(nda >= 6000, nda < 8000)] = 3
+            nda[np.logical_and(nda >= 8000, nda < 15000)] = 4
+        elif NumFlightLevels == 3:
+            nda[np.logical_and(nda >= 0, nda < 3000)] = 0
+            nda[np.logical_and(nda >= 3000, nda < 6000)] = 1
+            nda[np.logical_and(nda >= 6000, nda < 15000)] = 2
+
+        nda = tf.one_hot(nda, NumFlightLevels).numpy()
         nda = np.expand_dims(nda, axis=1)
         nda = np.expand_dims(nda, axis=1)