can't pickle tensorflow.python._tf_stack.StackSummary objects

I upgraded tensorflow to version 2.4.1 from tensorflow 2.2, and now I can no longer save my sklearn pipeline with a KerasRegressor inside. When I use joblib.dump, I get the following error message:

TypeError: can't pickle tensorflow.python._tf_stack.StackSummary objects

If I just try to pickle the file, I get this message:

TypeError: can't pickle weakref objects

If I use dill, I get the same error message as I do when I use joblib.

I've tried searching for this issue but I can't find much help. I've got no clue what the StackSummary object is

Here is a stand alone example that reproduces the issue:

import numpy as np

from keras.models import Sequential
from keras.layers import Dense
from keras import optimizers
from keras.wrappers.scikit_learn import KerasRegressor


from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.compose import TransformedTargetRegressor

from joblib import dump


def create_model(lr=10**-4, first_layer_nodes=100, amsgrad_yn=0,
                 second_layer_nodes=100, third_layer_nodes=100, loss='mse',
                 Leaky_alpha_slope=1, dropout_rate=0.01, output_shape=4, cal_points=[8,9]):

    model = Sequential()
    model.add(Dense(first_layer_nodes))
    Adam = optimizers.Adam(lr=lr, amsgrad=amsgrad_yn)
    model.add(Dense(output_shape, activation='linear'))
    model.compile(optimizer=Adam, loss=loss)
 
    return model


def make_pipeline(val_split=.2, verbose=1, callbacks=[], **parameters):
    NN = KerasRegressor(build_fn=create_model, validation_split=val_split, verbose=verbose, callbacks=callbacks, **parameters) 

    scaler = StandardScaler()
    y_scaler = StandardScaler()

    NN_model_with_y_scaling = TransformedTargetRegressor(regressor=NN, transformer=y_scaler)

    pipeline = Pipeline([
        ('preprocess', scaler),
        ('NN_model', NN_model_with_y_scaling)])

    return pipeline




model = make_pipeline(val_split=.2, verbose=0, **parameters)



x_train = np.random.rand(1000, 28)
y_train = np.random.rand(1000, 4)


model.fit(x_train, y_train)

dump(model, 'test.joblib')