Remove squeeze

Eden-Kramer-Lab · Sep 6, 2024 · 85e37e7 · 85e37e7
1 parent cb6ea90
commit 85e37e7
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/replay_trajectory_classification/decoder.py b/replay_trajectory_classification/decoder.py
@@ -326,9 +326,7 @@ def convert_results_to_xarray(
                 {
                     key: (
                         dims,
-                        mask(value, is_track_interior)
-                        .squeeze(axis=-1)
-                        .reshape(new_shape, order="F"),
+                        mask(value, is_track_interior).reshape(new_shape, order="F"),
                     )
                     for key, value in results.items()
                 },
@@ -340,9 +338,7 @@ def convert_results_to_xarray(
                 {
                     key: (
                         dims,
-                        mask(value, is_track_interior)
-                        .squeeze(axis=-1)
-                        .reshape(new_shape),
+                        mask(value, is_track_interior).reshape(new_shape),
                     )
                     for key, value in results.items()
                 },