Add retry handling to DeMemWM latent dataset

datasets/video/minecraft_video_dememwm_latent_dataset.py

@@ -90,6 +90,7 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
         self._target_offsets = np.arange(self.n_frames, dtype=np.int64) * self.frame_skip
         self._memory_segments = memory_segment_lengths(self.n_frames, self.memory_selection)
         self.memory_condition_length = sum(self._memory_segments[key] for key in SEGMENT_KEYS)
+        self.max_getitem_retries = int(cfg_get(cfg, "max_getitem_retries", 32))
         self._cached_arrays_path = None
         self._cached_arrays = None
         self.data_paths = self.get_data_paths(split)
@@ -159,7 +160,31 @@ class MinecraftVideoDeMemWMLatentDataset(torch.utils.data.Dataset):
         return file_idx, int(idx - prev)
 
     def __getitem__(self, idx: int):
-        return self.load_data(self.idx_remap[idx])
+        dataset_len = len(self)
+        if dataset_len == 0:
+            raise IndexError(
+                f"MinecraftVideoDeMemWMLatentDataset split {self.split!r} contains no clips under "
+                f"{self.save_dir / self.split}"
+            )
+
+        start_idx = int(idx)
+        errors = []
+        last_error = None
+        for attempt in range(max(1, self.max_getitem_retries)):
+            candidate_idx = (start_idx + attempt) % dataset_len
+            try:
+                return self.load_data(self.idx_remap[candidate_idx])
+            except Exception as exc:
+                last_error = exc
+                if len(errors) < 5:
+                    errors.append(f"idx={candidate_idx}: {type(exc).__name__}: {exc}")
+
+        details = "; ".join(errors) if errors else "no error details captured"
+        raise RuntimeError(
+            f"Failed to load MinecraftVideoDeMemWMLatentDataset item after "
+            f"{max(1, self.max_getitem_retries)} attempts starting at index {start_idx} "
+            f"for split {self.split!r}. Recent errors: {details}"
+        ) from last_error
 
     def load_data(self, idx: int):
         # === 1. Resolve clip and load arrays ===