Fix empty batch

dreem/datasets/base_dataset.py

@@ -110,6 +110,19 @@ def create_chunks(self) -> None:
 
                 self.label_idx = [self.label_idx[i] for i in sample_idx]
 
+            # check for batch with only a single element that correpsonds to the last frame of the video
+            remove_idx = None
+            for i, frame_chunk in enumerate(self.chunked_frame_idx):
+                if len(frame_chunk) == 1 and frame_chunk[0] % self.clip_length == 0:
+                    print(
+                        "Warning: Single frame batch; removing to avoid empty batch possibility with failed frame loading"
+                    )
+                    remove_idx = i
+                    break
+            if remove_idx is not None:
+                self.chunked_frame_idx.pop(remove_idx)
+                self.label_idx.pop(remove_idx)
+
         else:
             self.chunked_frame_idx = self.frame_idx
             self.label_idx = [i for i in range(len(self.labels))]
@@ -131,6 +144,7 @@ def no_batching_fn(self, batch: list[Frame]) -> list[Frame]:
         Returns:
             The batch
         """
+
         return batch
 
     def __getitem__(self, idx: int) -> list[Frame]:

dreem/inference/post_processing.py

@@ -160,7 +160,8 @@
         valid = dist.squeeze() < max_center_dist  # n_k x n_nonk
         # handle case where id_inds and valid is a single value
         # handle this better
-        if valid.ndim == 0: valid = valid.unsqueeze(0)
+        if valid.ndim == 0:
+            valid = valid.unsqueeze(0)
         if valid.ndim == 1:
             if id_inds.shape[0] == 1:
                 valid_mult = valid.float().unsqueeze(-1)