tensorflow: "Cannot take the length of shape with unknown rank" error

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04.3
• TensorFlow installed from (source or binary): pip
• TensorFlow version (use command below): tensorflow-addons==0.8.2 tensorflow-estimator==2.1.0 tensorflow-gpu==2.1.0 tensorflow-probability==0.9.0
• Python version: 3.6.9
• CUDA/cuDNN version: 10.1.243 / 7.6.5
• GPU model and memory: Tesla V100-SXM2 32 GB

Describe the current behavior I’m trying to switch from keras to tensorflow.keras (actually only to add TensorBoard callback to model.fit_generator and to be able to profile performance). Now I am getting

ValueError                                Traceback (most recent call last)
<ipython-input-6-70f0d0a1e94e> in <module>
…
~/Visual_Z2/ImageNet.py in _doLearning(self, epochCount, learningCallback, otherParams, initialEpochNum)
 #                             self.model.fit_generator
    430                                  validation_steps=testImageCount // self.batchSize,
    431                                  workers=2,
--> 432                                  verbose=2, callbacks=[tensorBoardCallback])
    433             #, summaryCallback])
    434             # Without make_one_shot_iterator - error fused convolution not supported

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/util/deprecation.py in new_func(*args, **kwargs)
    322               'in a future version' if date is None else ('after %s' % date),
    323               instructions)
--> 324       return func(*args, **kwargs)
    325     return tf_decorator.make_decorator(
    326         func, new_func, 'deprecated',
…
ValueError: in converted code:

    /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training_v2.py:677 map_fn
        batch_size=None)
    /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training.py:2469 _standardize_tensors
        exception_prefix='target')
    /usr/local/lib/python3.6/dist-packages/tensorflow_core/python/keras/engine/training_utils.py:529 standardize_input_data
…
ValueError: Cannot take the length of shape with unknown rank.

Describe the expected behavior Everything worked on keras with tensorflow as backend.

Other info / logs I read this can happen because of PyFunction, I really have it, but adding reshape/image.set_shape don’t fix the error. I build input dataset like this

def _loadTestImage(imageNum, label):
    imageData = self.imageDataset.getImage(imageNum, 'net', 'test')       # Returns numpy array (227, 227, 3) of floats
    return (imageData, keras.utils.to_categorical(label, num_classes=classCount))

def _tfLoadTestImage(imageNum, label):
    image, label = tf.py_function(_loadTestImage, [imageNum, label], [tf.float32, tf.int32]) 
    # Adding image.set_shape(crop_size + (3, ))
    # or tf.reshape(image, shape=crop_size + (3, ))
    # here doesn’t help
    return image, label

numDs = tf.data.Dataset.from_tensor_slices(imageNums)       # Shuffled indices of images [15, 22, 3, …] 
label_ds = tf.data.Dataset.from_tensor_slices(self.imageNumLabels)
ds = tf.data.Dataset.zip((numDs, label_ds))
ds = ds.repeat()
tfTestDataset = ds.shuffle(buffer_size=max(epochImageCount, 2000))
tfTestDataset = tfTestDataset.map(_tfLoadTestImage, num_parallel_calls=4)
tfTestDataset = tfTestDataset.batch(self.batchSize)

The problem looks identical to https://github.com/tensorflow/tensorflow/issues/24520#issuecomment-589119295 so I tried nightly build today (2.2.0-dev20200229) and still getting error, only

ValueError: in user code:

    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:499 train_function  *
        outputs = self.distribute_strategy.experimental_run_v2(
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:856 experimental_run_v2  **
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2112 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/distribute/distribute_lib.py:2470 _call_for_each_replica
        return fn(*args, **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py:471 train_step  **
        self.compiled_metrics.update_state(y, y_pred, sample_weight)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/compile_utils.py:348 update_state
        self._build(y_pred, y_true)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/compile_utils.py:267 _build
        self._metrics, y_true, y_pred)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/util/nest.py:1118 map_structure_up_to
        **kwargs)
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/util/nest.py:1214 map_structure_with_tuple_paths_up_to
        *flat_value_lists)]
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/util/nest.py:1213 <listcomp>
        results = [func(*args, **kwargs) for args in zip(flat_path_list,
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/util/nest.py:1116 <lambda>
        lambda _, *values: func(*values),  # Discards the path arg.
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/compile_utils.py:377 _get_metric_objects
        return [self._get_metric_object(m, y_t, y_p) for m in metrics]
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/compile_utils.py:377 <listcomp>
        return [self._get_metric_object(m, y_t, y_p) for m in metrics]
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/compile_utils.py:398 _get_metric_object
        y_t_rank = len(y_t.shape.as_list())
    /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/tensor_shape.py:1173 as_list
        raise ValueError("as_list() is not defined on an unknown TensorShape.")

    ValueError: as_list() is not defined on an unknown TensorShape.