Merge pull request #76 from fsx950223:master

PiperOrigin-RevId: 403124579
tensorflow · Oct 14, 2021 · 521b1e0 · 521b1e0
2 parents e5e245b + 4aefc6c
commit 521b1e0
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/tensorflow_estimator/python/estimator/head/base_head.py b/tensorflow_estimator/python/estimator/head/base_head.py
@@ -773,12 +773,11 @@ def all_class_ids(logits, n_classes):
 def all_classes(logits, n_classes, label_vocabulary=None):
   batch_size = tf.compat.v1.shape(logits)[0]
   if label_vocabulary:
-    classes_list = label_vocabulary
+    classes_list = tf.convert_to_tensor([label_vocabulary])
   else:
-    classes_list = tf.strings.as_string(tf.range(n_classes))
-  return tf.tile(
-      input=tf.compat.v1.expand_dims(input=classes_list, axis=0),
-      multiples=[batch_size, 1])
+    classes_list = tf.expand_dims(tf.range(n_classes), axis=0)
+    classes_list = tf.strings.as_string(classes_list)
+  return tf.tile(input=classes_list, multiples=[batch_size, 1])
 
 
 def classification_output(scores, n_classes, label_vocabulary=None):