Scikit-Learn wrapper for keras and RandomizedSearchCV result in infinite loop

Question

I tried some code examples from a ML book, but when I try to use the Scikit-learn wrapper for keras and want to tune the hyper parameters of my model with the RandomizedSearchCV, I somehow end up in an infinite loop. Any recommendations? My code looks like this:

import numpy as np
import tensorflow as tf
from tensorflow import keras 
from scipy.stats import reciprocal
from sklearn.model_selection import RandomizedSearchCV

fashion_mnist = keras.datasets.fashion_mnist
(X_train_full, y_train_full), (X_test, y_test) = fashion_mnist.load_data()
X_valid, X_train = X_train_full[:5000] / 255.0, X_train_full[5000:] / 255.0
y_valid, y_train = y_train_full[:5000], y_train_full[5000:]

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

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


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

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

The first training works but then it crashes, this is what the output looks like:

Any suggestions?