diff --git a/modules/deeplearning/unet_l1b_l2.py b/modules/deeplearning/unet_l1b_l2.py
index 22c3de10d19afd7c157c242cdc1cafd2a043e7e8..192f60fe7c9b93cfdae0679264e4b425fdcaf588 100644
--- a/modules/deeplearning/unet_l1b_l2.py
+++ b/modules/deeplearning/unet_l1b_l2.py
@@ -223,7 +223,7 @@ class UNET:
 
             f = self.train_label_files
             nda = np.load(f)
-            label = nda[idxs, 0, :, :]
+            label = nda[idxs, label_idx, :, :]
             label = np.expand_dims(label, axis=3)
         else:
             f = self.test_data_files