def convert_bn_to_train_mode(model):
  """Convert a trained model with batch norm layers to run in train mode."""
  bn_layers = [
      i for i, layer in enumerate(model.layers)
      if 'batch_normalization' in layer.name
  ]
  model_config = model.get_config()
  for i in bn_layers:
    model_config['layers'][i]['inbound_nodes'][0][0][-1]['training'] = True
  new_model = model.from_config(model_config)
  for i, layer in enumerate(new_model.layers):
    layer.set_weights(model.layers[i].get_weights())
  return new_model