4728eb68276af34969210e2054971e00f3fe7352,official/nlp/tasks/sentence_prediction.py,,predict,#Any#Any#Any#,232

Before Change

  dataset = orbit.utils.make_distributed_dataset(tf.distribute.get_strategy(),
                                                 task.build_inputs, params)
  // Set `num_steps` to -1 to exhaust the dataset.
  predictions = loop_fn(
      iter(dataset), num_steps=-1, state=[], reduce_fn=reduce_fn)
  return predictions

After Change

  dataset = orbit.utils.make_distributed_dataset(tf.distribute.get_strategy(),
                                                 task.build_inputs, params)
  outputs = utils.predict(predict_step, aggregate_fn, dataset)
  return outputs["predictions"]

Instances