diff --git a/modules/deeplearning/icing_cnn.py b/modules/deeplearning/icing_cnn.py
index 6ff325224f05c49a047fb4583bfc30484aff959b..7b20911736ece7f9f867ded0039de16abff63c05 100644
--- a/modules/deeplearning/icing_cnn.py
+++ b/modules/deeplearning/icing_cnn.py
@@ -206,15 +206,15 @@ class IcingIntensityNN:
         label = np.where(label == -1, 0, label)
 
         # binary, two class
+        #label = np.where(label != 0, 1, label)
+        #label = label.reshape((label.shape[0], 1))
+
+        keep = (label == 0) | (label == 3) | (label == 4) | (label == 5) | (label == 6)
+        data = data[keep,]
+        label = label[keep]
         label = np.where(label != 0, 1, label)
         label = label.reshape((label.shape[0], 1))
 
-        # keep = (label == 0) | (label == 3) | (label == 4) | (label == 5) | (label == 6)
-        # data = data[keep,]
-        # label = label[keep]
-        # label = np.where(label != 0, 1, label)
-        # label = label.reshape((label.shape[0], 1))
-
         if CACHE_DATA_IN_MEM:
             self.in_mem_data_cache[key] = (data, label)