Hi @User - I don't think it's possible out of the box via the YAML API. You can see the implementation here https://github.com/quantumblacklabs/kedro/blob/20f836695c2f1e72f262d1747e47b7b7352a4aa0/kedro/extras/datasets/tensorflow/tensorflow_model_dataset.py#L138 . If you want to do this, I think you would have to subclass the out of the box dataset and include this functionality yourself.