GridSearchCV with Keras

from sklearn.model_selection import GridSearchCV
from keras.wrappers.scikit_learn import KerasClassifier

def build_model():
  X = tf.keras.layers.Input(shape=[28, 28, 1])
  H = tf.keras.layers.Conv2D(64, kernel_size=5, padding='same', activation='swish')(X)
  H = tf.keras.layers.BatchNormalization()(H)

  H = tf.keras.layers.Conv2D(64, kernel_size=5, padding='same', activation='swish')(H)
  H = tf.keras.layers.MaxPool2D()(H)

  H = tf.keras.layers.Conv2D(64, kernel_size=3, padding='same', activation='swish')(H)
  H = tf.keras.layers.MaxPool2D()(H)

  H = tf.keras.layers.Conv2D(64, kernel_size=3, padding='same', activation='swish')(H)
  H = tf.keras.layers.MaxPool2D(pool_size=(2, 2), strides=(2, 2))(H)
  H = tf.keras.layers.Dropout(0.3)(H)

  H = tf.keras.layers.Flatten()(H)

  H = tf.keras.layers.Dense(128)(H)
  H = tf.keras.layers.BatchNormalization()(H)
  H = tf.keras.layers.Activation('swish')(H)

  H = tf.keras.layers.Dense(84)(H)
  H = tf.keras.layers.Activation('swish')(H)
  H = tf.keras.layers.Dropout(0.4)(H)

  Y = tf.keras.layers.Dense(10, activation='softmax')(H)
  model = tf.keras.models.Model(X, Y)
  model.summary()

  return model


estimator = KerasClassifier(build_fn = build_model, verbose=True, epochs=45, batch_size=64)

param_grid = {'batch_size': [32, 64], 'epochs': [10, 20, 30, 40]}
grid = GridSearchCV(estimator=estimator, param_grid=param_grid, cv=5, n_jobs=-1, verbose=4)
grid.fit(x_train, y_train)