From 8e6fd0885a7dca8e277431b06c9c1f7bad95f419 Mon Sep 17 00:00:00 2001 From: till-m Date: Tue, 23 Dec 2025 13:54:45 +0100 Subject: [PATCH] fix restricted trajectories --- the_well/data/datasets.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/the_well/data/datasets.py b/the_well/data/datasets.py index ccb6ab62..1254a9c0 100755 --- a/the_well/data/datasets.py +++ b/the_well/data/datasets.py @@ -389,8 +389,9 @@ def _build_restriction_set( - 1 ) if traj in trajectories_sampled: + n_windows = self.n_windows_per_trajectory[file_index] global_indices = global_indices + list( - range(0, self.n_windows_per_trajectory[file_index]) + range(current_index, current_index + n_windows) ) current_index += self.n_windows_per_trajectory[file_index] global_indices = np.array(global_indices)