diff --git a/src/mdevaluate/correlation.py b/src/mdevaluate/correlation.py
index 3b67777..7d2eee6 100644
--- a/src/mdevaluate/correlation.py
+++ b/src/mdevaluate/correlation.py
@@ -99,9 +99,9 @@ def shifted_correlation(
             return get_correlation(frames, start_frame, index, shifted_idx)
         else:
             index = selector(frames[start_frame])
-            if len(index.shape) == 1:
+            if index.ndim == 1:
                 return get_correlation(frames, start_frame, index, shifted_idx)
-            elif len(index.shape) == 2:
+            elif index.ndim == 2:
                 correlations = []
                 for ind in index:
                     correlations.append(
@@ -110,7 +110,7 @@ def shifted_correlation(
                 return correlations
             else:
                 raise ValueError(
-                    f"Index list of selector has {len(index.shape)} dimensions, "
+                    f"Index list of selector has {index.ndim} dimensions, "
                     "but should have 1 or 2"
                 )