resolved #336

neuronets · Jun 10, 2024 · f0396cb · f0396cb
1 parent a8e960b
commit f0396cb
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/nobrainer/dataset.py b/nobrainer/dataset.py
@@ -146,6 +146,8 @@ def from_tfrecords(
             ds_obj.map_labels(
                 label_mapping=label_mapping, num_parallel_calls=num_parallel_calls
             )
+
+        ds_obj.filter_zero_volumes()
         # TODO automatically determine batch size
         ds_obj.batch(1)
 
@@ -385,3 +387,9 @@ def repeat(self, n_repeats):
         # through once.
         self.dataset = self.dataset.repeat(n_repeats)
         return self
+
+    def filter_zero_volumes(self):
+        self.dataset = self.dataset.filter(
+            lambda x, y: tf.cast(tf.math.reduce_sum(y), dtype="bool")
+        )
+        return self