diff --git a/modules/deeplearning/icing_cnn.py b/modules/deeplearning/icing_cnn.py
index a65f86d163d81c968f8985d4b8604d1d1e36f398..0c7add6d0087c13e5db5740672b7383e52ee9287 100644
--- a/modules/deeplearning/icing_cnn.py
+++ b/modules/deeplearning/icing_cnn.py
@@ -16,8 +16,13 @@ CACHE_DATA_IN_MEM = False
 
 PROC_BATCH_SIZE = 4096
 PROC_BATCH_BUFFER_SIZE = 50000
+
 NumClasses = 2
-NumLogits = 1
+if NumClasses == 2:
+    NumLogits = 1
+else:
+    NumLogits = NumClasses
+
 BATCH_SIZE = 128
 NUM_EPOCHS = 100