Update _with_tensor_ranks_only TODOs
@@ -817,6 +817,7 @@ class GeneratorDataAdapter(DataAdapter):
self._first_batch_size = int(tf.nest.flatten(peek)[0].shape[0])
def _get_tensor_spec(t):
+ # TODO(b/226395276): Remove _with_tensor_ranks_only usage.
return type_spec.type_spec_from_value(t)._with_tensor_ranks_only() # pylint: disable=protected-access
output_signature = tf.nest.map_structure(_get_tensor_spec, peek)
