title: Unable to interpret an argument of type tensorflow.python.data.ops.dataset_ops.PrefetchDataset as a TFF value in iterative process
body: <p>I'm trying to run a classification simulation in tff, but I'm getting this error:</p>
  <pre><code>TypeError: Unable to interpret an argument of type tensorflow.python.data.ops.dataset_ops.PrefetchDataset as a TFF value.
  </code></pre>
  <p>Here is the code I'm using</p>
  <pre><code>client_lr = 1e-3
  server_lr = 1e-1
  NUM_ROUNDS = 200
  NUM_EPOCHS = 5
  BATCH_SIZE = 2048
  EPOCHS = 400
  TH = 0.5
  
  def base_model():
    return Sequential([
          Dense(256, activation='relu', input_shape=(x_train.shape[-1],)),
          Dropout(0.5),
          Dense(256, activation='relu'),
          Dropout(0.5),
          Dense(256, activation='relu'),
          Dropout(0.5),
          Dense(1, activation='sigmoid'),
      ])
  
  client_train_dataset = collections.OrderedDict()
  for i in range(1, total_clients+1):
    client_name = &quot;client_&quot; + str(i)
    start = samples_per_set * (i-1)
    end = samples_per_set * i
    data = collections.OrderedDict((('y', y_train[start:end]), ('x', x_train[start:end])))
    client_train_dataset[client_name] = data
  
  train_dataset = tff.simulation.FromTensorSlicesClientData(client_train_dataset)
  
  sample_dataset = train_dataset.create_tf_dataset_for_client(train_dataset.client_ids[0])
  sample_element = next(iter(sample_dataset))
  PREFETCH_BUFFER = 10
  SHUFFLE_BUFFER = samples_per_set
  
  def preprocess(dataset):
  
    def batch_format_fn(element):
      return collections.OrderedDict(
          x=element['x'],
          y=tf.reshape(element['y'], [-1, 1]))
  
    return dataset.repeat(NUM_EPOCHS).shuffle(SHUFFLE_BUFFER).batch(BATCH_SIZE).map(batch_format_fn).prefetch(PREFETCH_BUFFER)
  
  preprocessed_sample_dataset = preprocess(sample_dataset)
  sample_batch = tf.nest.map_structure(lambda x: x.numpy(), next(iter(preprocessed_sample_dataset)))
  
  def make_federated_data(client_data, client_ids):
      return [preprocess(client_data.create_tf_dataset_for_client(x)) for x in client_ids]
  
  federated_train_data = make_federated_data(train_dataset, train_dataset.client_ids)
  
  def model_tff():
    model = base_model()
    return tff.learning.from_keras_model(
        model,
        input_spec=preprocessed_sample_dataset.element_spec,
        loss=tf.keras.losses.BinaryCrossentropy(),
        metrics=[
                 tfa.metrics.F1Score(num_classes=1, threshold=TH),
                 keras.metrics.Precision(name=&quot;precision&quot;, thresholds=TH),
                 keras.metrics.Recall(name=&quot;recall&quot;, thresholds=TH)
                ])
  
  iterative_process = tff.learning.build_federated_averaging_process(
      model_tff,
      client_optimizer_fn=lambda: optimizers.Adam(learning_rate=client_lr),
      server_optimizer_fn=lambda: optimizers.SGD(learning_rate=server_lr))
  
  state = iterative_process.initialize()
  
  federated_model = None
  
  for round_num in range(1, NUM_ROUNDS+1):
      state, tff_metrics = iterative_process.next(state, federated_train_data) # THE ERROR IS HERE
      federated_model = base_model()
      federated_model.compile(optimizer=optimizers.Adam(learning_rate=client_lr),
                          loss=tf.keras.losses.BinaryCrossentropy(),
                          metrics=[
                                tfa.metrics.F1Score(num_classes=1, threshold=TH),
                                keras.metrics.Precision(name=&quot;precision&quot;, thresholds=TH),
                                keras.metrics.Recall(name=&quot;recall&quot;, thresholds=TH)
                                ])
      state.model.assign_weights_to(model=federated_model)
      federated_result = federated_model.evaluate(x_val, y_val, verbose=1, return_dict=True)
      
  federated_test = federated_model.evaluate(x_test, y_test, verbose=1, return_dict=True)
  </code></pre>
  <p>I'm using this creditcard dataset: <a href="https://www.kaggle.com/mlg-ulb/creditcardfraud" rel="nofollow noreferrer">https://www.kaggle.com/mlg-ulb/creditcardfraud</a></p>
  <p>The <code>federated_train_data</code> is a list of <code>&lt;PrefetchDataset shapes: OrderedDict([(x, (None, 29)), (y, (None, 1))]), types: OrderedDict([(x, tf.float64), (y, tf.int64)])&gt;</code>, just like the tutorial from the Tensorflow Federated website <a href="https://www.tensorflow.org/federated/tutorials/federated_learning_for_image_classification" rel="nofollow noreferrer">Federated Learning for Image Classification</a>.</p>