Fixes a test failure that was caused by the numba update in 4362401fa377311b1fbbdb7bda2abc406be2589d.