diff --git a/modules/deeplearning/icing.py b/modules/deeplearning/icing.py
index 6c3ee9026f704fb7e8cce916e2b7cdd8a0635ec9..5e67a711cc7e8ae11adebe27909f125e9ae81d6c 100644
--- a/modules/deeplearning/icing.py
+++ b/modules/deeplearning/icing.py
@@ -426,7 +426,7 @@ class IcingIntensityNN:
 
     @tf.function
     def train_step(self, mini_batch):
-        inputs = [mini_batch[0], mini_batch[1], mini_batch[2]]
+        inputs = [mini_batch[0], mini_batch[1]]
         labels = mini_batch[2]
         with tf.GradientTape() as tape:
             pred = self.model(inputs, training=True)