MirroredStrategy: Inputs to eager execution function cannot be Keras symbolic tensors

I'm using tensorflow 2.0 to train a model, my data pipeline and model train fine with tf.keras.Model.fit() on a single gpu, but when I try running it on multi gpu with MirroredStrategy.

My dataset is obtained with something like this:

dataset = tf.data.TFRecordDataset(records)
dataset.map(parse_and_decode)
map_augmentation = lambda *args: tf.py_function(
    augmentation,
    [*args],
    4*[tf.float32]
)
dataset.map(map_augmentation)

def map_ground_truth(arg1, arg2, arg3, arg4):
    ret_values = tf.py_function(
        get_ground_truth,
        [arg1, arg2, arg3, arg4],
        [tf.float32, tf.float32, tf.float32, tf.float32, tf.float32, tf.float32
    )
    # set_shape based on https://github.com/tensorflow/tensorflow/issues/24520
    for ret_val in ret_values:
        ret_val.set_shape([None for _ in range(3)])
    return *ret_values

dataset.map(map_ground_truth)

model.fit(dataset)

This works ok without using MirroredStrategy but when I introduce it I get this error:

_SymbolicException: Inputs to eager execution function cannot be Keras symbolic tensors, but found [<tf.Tensor 'crowd_mask:0' shape=(1, 401, 401, 1) dtype=float32>, <tf.Tensor 'unannotated_mask:0' shape=(1, 401, 401, 1) dtype=float32>, <tf.Tensor 'kp_maps_true:0' shape=(1, 401, 401, 17) dtype=float32>, <tf.Tensor 'mid_offsets/Identity:0' shape=(1, 13, 13, 64) dtype=float32>]

(in the strategy.scope() are the model and losses creations, the model compile and the fit. No manual call to strategy.experimental_distribute_dataset is done on the dataset as this is aparently handled by the fit function).