diff --git a/modules/deeplearning/icing_fcn.py b/modules/deeplearning/icing_fcn.py
index 6c669f7ae6df88c29803c676e337ef45f723c681..b4b91fdd4dc13d398555f9444ee02ffd3a4984b3 100644
--- a/modules/deeplearning/icing_fcn.py
+++ b/modules/deeplearning/icing_fcn.py
@@ -344,6 +344,8 @@ class IcingIntensityFCN:
         nda[np.logical_and(nda >= 8000, nda < 15000)] = 4
 
         nda = tf.one_hot(nda, 5).numpy()
+        nda = np.expand_dims(nda, axis=1)
+        nda = np.expand_dims(nda, axis=1)
 
         return nda
 
@@ -903,9 +905,7 @@ class IcingIntensityFCN:
     def build_model(self):
         cnn = self.build_cnn()
         if self.USE_FLIGHT_ALTITUDE:
-            fa = np.expand_dims(self.inputs[1], axis=1)
-            fa = np.expand_dims(fa, axis=1)
-            cnn = tf.keras.layers.concatenate([cnn, fa])
+            cnn = tf.keras.layers.concatenate([cnn, self.inputs[1]])
         self.build_fcl(cnn)
         self.model = tf.keras.Model(self.inputs, self.logits)