diff --git a/python/tests/test_imputation.py b/python/tests/test_imputation.py
index cbffde514f..dd2a97986c 100644
--- a/python/tests/test_imputation.py
+++ b/python/tests/test_imputation.py
@@ -484,8 +484,7 @@ def test_compare_imputed_alleles(input_ref, input_query, expected):
         imputed_alleles[i], _ = tests.beagle.run_beagle(
             input_ref, input_query[i], pos, miscall_rate=0.0001, ne=10.0
         )
-        num_diff = np.sum(imputed_alleles[i] != expected_subset[i])
-        assert num_diff == 0, f"Wrongly imputed alleles: {num_diff}"
+        np.testing.assert_array_equal(imputed_alleles[i], expected_subset[i])
 
 
 @pytest.mark.parametrize(