# Define the Keras model without any Horovod-specific parameters
model = Sequential()
model.add(Conv2D(32, kernel_size=(3, 3),
                 activation='relu',
                 input_shape=(28, 28, 1)))
model.add(Conv2D(64, (3, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Dropout(0.25))
model.add(Flatten())
model.add(Dense(128, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(10, activation='softmax'))
 
optimizer = keras.optimizers.Adadelta(1.0)
loss = keras.losses.categorical_crossentropy

# Train a Horovod Spark Estimator on the DataFrame
keras_estimator = hvd.KerasEstimator(
  num_proc=num_proc,
  store=store,
  model=model,
  optimizer=optimizer,
  loss=loss,
  metrics=['accuracy'],
  feature_cols=['features'],
  label_cols=['label_vec'],
  batch_size=batch_size,
  epochs=epochs,
  verbose=1)
 
keras_model = keras_estimator.fit(train_df).setOutputCols(['label_prob'])