How to fix numpy.random.choice output nested inside a for loop when importing from python file?

Problem:

I am trying to generate multiple lists of random numbers, of differing length. Inside a for loop of length num_baskets, I am using np.random.choice to generate a number for the length of each successive list n. Then, again, np.random.choice generates a list of random numbers of length n.

Expected output:

List of length num_baskets, with each item in the list another list of random length n.

Current output: List of length num_baskets but with uniform sublist lengths. However, each time I run the function, sublist length is different (but still uniform) per call.

What I have tried:

When I import the function from a python file, (e.g. from python_file import create_baskets there is deviation from expected output. Despite redefining n in each loop, all the output lists have the same length.

However, when I copy and paste the function and define it inside a jupyter notebook, I get the expected output with different list lengths.

My code:

import numpy as np

def create_baskets(num_baskets, max_basket_size, unique_items):
    """
    Create list of baskets of variable length >= 3

    Parameters
    ----------
    num_baskets: 
        number of baskets (sub-lists)
    max_basket_size: maximum basket size
        Baskets will be of size range(3, max_basket_size)
    unique_items:
        number of unique items

    Returns
    -------
    ret: list of "baskets", of randomly generated length
    """
    baskets = []
    for i in range(num_baskets):
        n = np.random.choice(range(3, max_basket_size), 1)
        basket = np.random.choice(range(0, unique_items), n, replace=False)
        baskets.append(basket)
    return baskets

I am not sure if there is something that fundamentally wrong/improveable with how the function is written or if it is a problem with the import.