diff --git a/python/tests/beagle.py b/python/tests/beagle.py
index 46cb59a4c8..a3ecc08e34 100644
--- a/python/tests/beagle.py
+++ b/python/tests/beagle.py
@@ -128,13 +128,13 @@ def get_switch_prob(pos, h, ne=1e6):
 
     :param numpy.ndarray pos: Site positions.
     :param int h: Number of reference haplotypes.
-    :param float ne: Effective population size.
+    :param int ne: Effective population size.
     :return: Switch probabilities.
     :rtype: numpy.ndarray
     """
     assert pos[0] > 0, "Site positions are not 1-based."
     assert isinstance(h, int), "Number of reference haplotypes is not an integer."
-    assert isinstance(ne, float), "Effective population size is not a float."
+    assert isinstance(ne, int), "Effective population size is not an integer."
     # Get genetic distance between adjacent markers.
     cm = np.concatenate([[0], convert_to_genetic_map_position(pos)])
     gen_dist = cm[1:] - cm[:-1]