python multiprocessing pool.map hangs when calling tensorflow/keras model - tensorflow

Question

I use pool.map from multiprocessing to parallelize my python code. When I call my tensorflow/keras model with pool.map, the code hangs if my neural network is larger than a certain size. I still have plenty of RAM available, and calling the model outside of pool works fine.
I use python 3.7, tensorflow 2.3 on linux.
A mwe is provided below, it is also on colab:
def my_function(i):
a = MODEL(np.array(i).reshape(1,1))
print('foo', i)
return a
THREADS = os.cpu_count()
N = 4
NEURONS = 150000 # works for 100000, hangs for 150000
MODEL = tf.keras.Sequential([tf.keras.layers.Dense(NEURONS, input_shape=(1,))])
my_function(10) # works fine
pool = multiprocessing.Pool(THREADS)
_ = pool.map(my_function, range(N)) # hangs
pool.close()
pool.join()
Any idea what the issue is? How can I call a large model in parallel?
Edit: the size of a is not the issue, and the code hangs only if tf.keras is called once outside of pool, see mwe below and colab. The critical number of neurons is lower than in the original example. Any idea?
def my_function(i):
print('start', i)
model = tf.keras.Sequential([tf.keras.layers.Dense(NEURONS, input_shape=(1,))])
print('finish', i)
return None
THREADS = os.cpu_count()
N = 4
NEURONS = 20000 # works with 10000, not with 20000
# works
pool = multiprocessing.Pool(THREADS)
_ = pool.map(my_function, range(N))
pool.close()
pool.join()
# works
my_function(10)
# doesn't work if many neurons
pool = multiprocessing.Pool(THREADS)
_ = pool.map(my_function, range(N))
pool.close()
pool.join()