def METHOD_NAME(optimizer, initial_weights):
  # Clear Keras session to reset device assignment
  keras.backend._SESSION.session = None
  strategy = strategy_cls()
  with strategy.scope():
    train_ds = get_input_datasets()
    model = get_simple_bias_model()
    model.set_weights(initial_weights)
    model.compile(loss='mae', optimizer=optimizer, metrics=['mae'])
  return {
      'trained_loss_and_accuracy':
          model.fit(x=train_ds, epochs=20).history,
      'trained_weights':
          model.get_weights(),
  }