Scikit-Learn wrapper and RandomizedSearchCV: RuntimeError

I am reading the book

"Hands-On Machine Learning with Scikit-Learn, Keras, and Tensorflow: Concepts, Tools, and Techniques to Build Intelligent Systems"

and in the Chapter 11 (Introduction to ANN with Keras) is explained that one can wrap a tensorflow model in scikit-learn to use some useful tools, like RandomizedSearchCVwhich is quite useful for random search of ANN hyperparameters (ANN structure, learning rate, activation functions, etc)

But I get a strange error at the end of the Randomized Search. Specifically, after the random search, at the end of every combinations I get this:

RuntimeError                              Traceback (most recent call last)
<ipython-input-35-094d1018c18c> in <module>()
     13 rnd_search_cv.fit(X_train, y_train, epochs=100, 
     14                   validation_data=(X_valid, y_valid),
---> 15                   callbacks=[keras.callbacks.EarlyStopping(patience=10)])

1 frames
/usr/local/lib/python3.6/dist-packages/sklearn/model_selection/_search.py in fit(self, X, y, groups, **fit_params)
    734             # of the params are estimators as well.
    735             self.best_estimator_ = clone(clone(base_estimator).set_params(
--> 736                 **self.best_params_))
    737             refit_start_time = time.time()
    738             if y is not None:

/usr/local/lib/python3.6/dist-packages/sklearn/base.py in clone(estimator, safe)
     80             raise RuntimeError('Cannot clone object %s, as the constructor '
     81                                'either does not set or modifies parameter %s' %
---> 82                                (estimator, name))
     83     return new_object
     84 

RuntimeError: Cannot clone object <tensorflow.python.keras.wrappers.scikit_learn.KerasRegressor object at 0x7f16ce468fd0>, as the constructor either does not set or modifies parameter learning_rate

I followed every step in the chapter, namely:

Function for model parameterization

# build model given a set of parameters
input_shape = X_train[0].shape
X_new = X_test[:3]

def build_model(n_hidden=1, n_neurons=30, learning_rate=3e-3):
    model = keras.models.Sequential()
    model.add(keras.layers.InputLayer(input_shape=input_shape))
    for layer in range(n_hidden):
        model.add(keras.layers.Dense(n_neurons, activation="relu"))
    model.add(keras.layers.Dense(1))
    optimizer = keras.optimizers.SGD(lr=learning_rate)
    model.compile(optimizer=optimizer, loss="mse")
    return model

Scikit model wrapper

keras_reg = keras.wrappers.scikit_learn.KerasRegressor(build_model)

I also tested the model, and it worked just fine

keras_reg.fit(X_train, y_train, epochs=100, validation_data=(X_valid, y_valid), 
              callbacks=[keras.callbacks.EarlyStopping(patience=10)])
mse_test = keras_reg.score(X_test, y_test)
y_pred = keras_reg.predict(X_new)

But when I used the RandomizedSearchCV

# use RandomSearch (or grid search)

from scipy.stats import reciprocal
from sklearn.model_selection import RandomizedSearchCV

param_distribs = {
    "n_hidden": [0, 1, 2, 3],
    "n_neurons": np.arange(1, 100),
    "learning_rate": reciprocal(3e-4, 3e-2)
}

rnd_search_cv = RandomizedSearchCV(keras_reg, param_distribs, n_iter=10, cv=3)
rnd_search_cv.fit(X_train, y_train, epochs=100, 
                  validation_data=(X_valid, y_valid),
                  callbacks=[keras.callbacks.EarlyStopping(patience=10)])

I get the above RuntimeError.

I am working on colab, with tensorflow 2.3.0

import tensorflow as tf
tf.__version__

2.3.0

Does someone know why?