dataset: Add option to only catch NoneChucksSkipExceptions

nonechucks/dataset.py

@@ -2,17 +2,19 @@
 import torch.utils.data
 
 from nonechucks.utils import memoize
+from nonechucks.utils import NoneChucksSkipException
 
 
 class SafeDataset(torch.utils.data.Dataset):
     """A wrapper around a torch.utils.data.Dataset that allows dropping
     samples dynamically.
     """
 
-    def __init__(self, dataset, eager_eval=False):
+    def __init__(self, dataset, eager_eval=False, catch_all = True):
         """Creates a `SafeDataset` wrapper around `dataset`."""
         self.dataset = dataset
         self.eager_eval = eager_eval
+        self.catch_all = True
         # These will contain indices over the original dataset. The indices of
         # the safe samples will go into _safe_indices and similarly for unsafe
         # samples.
@@ -41,6 +43,9 @@ def _safe_get_item(self, idx):
                 self._safe_indices.append(idx)
             return sample
         except Exception as e:
+            if not self.catch_all:
+                if not isinstance(e, NoneChucksSkipException):
+                    raise
             if isinstance(e, IndexError):
                 if invalid_idx:
                     raise