Change indexing to enable 0 drop_steps.

simonpf · Sep 19, 2024 · e2b37be · e2b37be
1 parent 87e5a37
commit e2b37be
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/chimp/testing.py b/chimp/testing.py
@@ -189,13 +189,15 @@ def process_tile(
                 y_preds_k_r = y_preds_k
                 targets_k_r = targets_k
             input_maps_r = input_maps
+            age_maps_r = age_maps
 
             if drop_steps is not None:
-                y_preds_k_r = y_preds_k_r[drop_steps:-drop_steps]
-                targets_k_r = targets_k_r[drop_steps:-drop_steps]
-                input_maps_r = input_maps[drop_steps:-drop_steps]
-                age_maps_r = age_maps[drop_steps:-drop_steps]
-
+                left = drop_steps
+                right = len(y_preds_k_r) - drop_steps
+                y_preds_k_r = y_preds_k_r[left:right]
+                targets_k_r = targets_k_r[left:right]
+                input_maps_r = input_maps[left:right]
+                age_maps_r = age_maps[left:right]
 
             # Evaluate retrieval.
             for step, (y_pred_k, target_k, input_map) in enumerate(zip(