diff --git a/src/agora/utils/association.py b/src/agora/utils/association.py
index 13b9cbe408e608b346a1fa91d7704f7739d4e846..656aa787fbf4029f964bfc4ba971b80b2f1177e4 100644
--- a/src/agora/utils/association.py
+++ b/src/agora/utils/association.py
@@ -71,7 +71,7 @@ def validate_association(
     if association.ndim < 3:
         # Reshape into 3-D array for broadcasting if neded
         association = np.stack(
-            (association[:, [0, 1]], association[:, [0, 2]]), axis=2
+            (association[:, [0, 1]], association[:, [0, 2]]), axis=1
         )
 
     # Compare existing association with available indices