def update_param_grid(param_grid, sizes):
  """Update parameters of the grid search.

  Args:
    param_grid: current grid search dictionary
    sizes: dataset.shape
  Returns:
    param_grid: updated grid search dictionary
  """
  if 'h_units' in param_grid.keys():
    # NN and simple_linear data
    if sizes[1] > 1:
      param_grid['h_units'] = [i for i in param_grid['h_units'] if i < sizes[1]]
    else:
      param_grid['h_units'] = [1]

    if sizes[0] > 10001:
      param_grid['batch_size'] = [600]
  return param_grid