Add tensor shape relaxation back in
@@ -817,7 +817,7 @@ class GeneratorDataAdapter(DataAdapter):
self._first_batch_size = int(tf.nest.flatten(peek)[0].shape[0])
def _get_tensor_spec(t):
- return type_spec.type_spec_from_value(t)
+ return type_spec.type_spec_from_value(t)._with_tensor_ranks_only()
output_signature = tf.nest.map_structure(_get_tensor_spec, peek)
