Numba function fails when imported from another script, but works well when compiled manually

Here is a reproducible example. The numba function "is_in_set_pnb" doesn't work when is imported from "functions.py" script. However, it works just fine when is defined in the same script.

functions.py script

import numpy as np
import numba as nb
import pandas as pd

# corrFilter function
def corrFilter(df, threshold):
    corr_matrix = df.corr().abs()
    flat_matrix = (corr_matrix.where(np.triu(np.ones(corr_matrix.shape), k=1).astype(np.bool))
                 .stack()
                 .sort_values(ascending=False))
    flat_matrix = pd.DataFrame(flat_matrix).reset_index()
    flat_matrix.columns=["var1", "var2", "correlation"]
    flat_matrix = flat_matrix.reindex(flat_matrix.correlation.sort_values(ascending=False).index).reset_index().drop(["index"],axis=1)
    filtered_matrix = flat_matrix[flat_matrix.correlation > threshold]
    pairs_to_remove = filtered_matrix[["var1", "var2"]].to_numpy()
    
    return pairs_to_remove

# Numba function / numpy.isin() improved version
@nb.jit(parallel=True)
def is_in_set_pnb(a, b):
    shape = a.shape
    a = a.ravel()
    n = len(a)
    result = np.full(n, False)
    set_b = set(b)
    for i in nb.prange(n):
        if a[i] in set_b:
            result[i] = True
    return result.reshape(shape)

main.py script

from sklearn import datasets
import pandas as pd
import numpy as np
import itertools
from functions import *

# Loading data
iris = datasets.load_iris()
data = iris.data
data = pd.DataFrame(data)
data.columns = ["var1", "var2", "var3", "var4"]

# Get pairs with a correlation higher than 0.5
remove_pairs = corrFilter(data, 0.5)
remove_pairs_check = np.sum(remove_pairs, axis=1)
# Remove triplets with pairs with a higher correlation than 0.5
triplets = [list(k) for k in itertools.combinations(data.columns, 3)]
array_triplets = np.vstack(triplets).astype(object)
n, d = array_triplets.shape
pair1 = np.sum(array_triplets[:,[0,1]], axis=1).reshape(n, 1)
pair2 = np.sum(array_triplets[:,[1,2]], axis=1).reshape(n, 1)
pair3 = np.sum(array_triplets[:,[0,2]], axis=1).reshape(n, 1)
array_triplets_check = np.concatenate((pair1, pair2, pair3), axis =1)

# Check if each pair is in the remove_pairs_check list
array_triplets_check_v1 = np.in1d(array_triplets_check, remove_pairs_check).reshape(n, d)
# Using numba function from "functions.py" script"
# This fails!!
array_triplets_check_v2 = is_in_set_pnb(array_triplets_check, remove_pairs_check)

However, if we define the numba function within the main.py script:

import numba as nb
@nb.jit(parallel=True)
def is_in_set_pnb(a, b):
    shape = a.shape
    a = a.ravel()
    n = len(a)
    result = np.full(n, False)
    set_b = set(b)
    for i in nb.prange(n):
        if a[i] in set_b:
            result[i] = True
    return result.reshape(shape)

# This works properly!!!
array_triplets_check_v2 = is_in_set_pnb(array_triplets_check, remove_pairs_check)

What I am missing?? I am using:

numba                     0.50.1           py38h47e9c7a_0