Allow Keras remove_squeezable_dimensions to accept all composite tensors.
@@ -111,9 +111,10 @@ def remove_squeezable_dimensions(
Tuple of `labels` and `predictions`, possibly with last dim squeezed.
"""
with backend.name_scope(name or 'remove_squeezable_dimensions'):
- if not isinstance(predictions, tf.RaggedTensor):
+ # tf.is_tensor returns True if predictions is a tensor or composite tensor.
+ if not tf.is_tensor(predictions):
predictions = tf.convert_to_tensor(predictions)
- if not isinstance(labels, tf.RaggedTensor):
+ if not tf.is_tensor(labels):
labels = tf.convert_to_tensor(labels)
predictions_shape = predictions.shape
predictions_rank = predictions_shape.ndims
