New RunLLM response generated: >>>*# Question* I am experimenting with upgrading my workflow to use

tensorflow=2.16.2

. This includes some breaking changes with how

keras

works. It's no longer bundled as

tensorflow.keras

and is instead it's own package. It also has some breaking changes about how models are saved and loaded. Details in this migration guide. I'm getting an error in a flyte task that returns a trained model. I think this is caused by the TensorFlowModelTransformer using the keras 2 style. To make this work with keras 3, I think the fix would be something like this:

-        local_path = ctx.file_access.get_random_local_path()
+        local_path = ctx.file_access.get_random_local_path() + ".keras"
        pathlib.Path(local_path).parent.mkdir(parents=True, exist_ok=True)

        # save model in SavedModel format
-        tf.keras.models.save_model(python_val, local_path)
+        keras.saving.save_model(python_val, local_path)

This is untested code BTW. And it would require similar changes in loading the model. It might require a new transformer for Keras 3 specifically. I didn't see anything related to this in the github issues. I'm wondering if others have had this problem as well. # Answer The issue you're encountering is likely due to the changes in how Keras models are saved and loaded in Keras 3, as you mentioned. The

TensorFlowModelTransformer

in Flytekit is indeed designed for Keras 2, and the migration to Keras 3 requires adjustments in the save and load methods. According to the migration guide, Keras 3 has introduced breaking changes, including how models are saved and loaded. Your proposed changes to the

TensorFlowModelTransformer

seem appropriate, but they need to be tested and verified. Would you like more detailed steps on how to implement and test these changes, or further assistance with a specific part of the migration? Sources: - https://github.com/flyteorg/flyte/issues/2570 - https://github.com/flyteorg/flyte/issues/5706 Note: I wasn’t able to find highly relevant data sources, but above are a few potentially relevant links. Slack Link https://flyte-org.slack.com/archives/CP2HDHKE1/p1726160014.882229 👍 0 button 👎 0 button 💡 Hint: Mention @runllm in the thread for followups.