Fix numba depreciation warning