dataset make sure n_train _val at least 1

profjsb · Oct 5, 2019 · b9bb1a2 · b9bb1a2
1 parent ca97811
commit b9bb1a2
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/deepCR/dataset.py b/deepCR/dataset.py
@@ -144,9 +144,9 @@ def __init__(self, image, mask, ignore=None, sky=None, aug_sky=[0, 0], part=None
             ignore = np.zeros_like(image)
 
         if part == 'train':
-            s = np.s_[:int(len * f_train)]
+            s = np.s_[:max(1, int(len * f_train))]
         elif part == 'val':
-            s = np.s_[int(len * f_train):]
+            s = np.s_[min(len - 1, int(len * f_train)):]
         else:
             s = np.s_[0:]