At inference time to reduce the computational cost I can set the maximum number of geometries to, for e.g., 10000: distances, indices = kdtree.query(points_coords_array, k=min(10000,len(geoseries_geometries_epsg)))