diff --git a/python/tests/beagle.py b/python/tests/beagle.py
index d197150c6c..170984e982 100644
--- a/python/tests/beagle.py
+++ b/python/tests/beagle.py
@@ -528,7 +528,7 @@ def get_map_alleles(allele_probs):
     """
     assert not np.any(allele_probs < 0), "Allele probabilities have negative values."
     assert not np.any(np.isnan(allele_probs)), "Allele probabilities have NaN values."
-    x = len(allele_probs)
+    x = allele_probs.shape[0]
     imputed_alleles = np.zeros(x, dtype=int)
     # TODO: Vectorise over the imputed markers
     for i in np.arange(x):