Skip to content

GitLab

Commit 10643806 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 326722720
parent 6e3b5a08
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
official/nlp/tasks/sentence_prediction_test.py
View file @ 10643806
...@@ -248,6 +248,9 @@ class SentencePredictionTaskTest(tf.test.TestCase, parameterized.TestCase): ...@@ -248,6 +248,9 @@ class SentencePredictionTaskTest(tf.test.TestCase, parameterized.TestCase):
predictions = sentence_prediction.predict(task, test_data_config, model) predictions = sentence_prediction.predict(task, test_data_config, model)
self.assertLen(predictions, num_examples) self.assertLen(predictions, num_examples)
for prediction in predictions:
self.assertEqual(prediction.dtype,
tf.int64 if num_classes > 1 else tf.float32)
if __name__ == "__main__": if __name__ == "__main__":
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment