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

Before Change

      state.extend(per_replica_batch_predictions)
    return state

  loop_fn = orbit.utils.create_loop_fn(predict_step)
  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(

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