I really need your help with this cryptic tensorflow behaviour. Let me explain the scenario.

1. Train model (using gradient tape)
2. model.save(my_dir)
3. tf.keras.models.load_model(my_dir, compile=False)
4. retrain model.
5. model.save(my_dir1)
6. tf.keras.models.load_model(my_dir1, compile=False) -> AttributeError

StackTrace:

self.model = tf.keras.models.load_model(my_dir1, compile=False)
  File "/usr/local/lib/python3.8/dist-packages/keras/utils/traceback_utils.py", line 70, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/usr/local/lib/python3.8/dist-packages/keras/saving/legacy/saved_model/load.py", line 625, in _revive_layer_or_model_from_config
    obj._set_dtype_policy(metadata["dtype"])
AttributeError: 'Ones' object has no attribute '_set_dtype_policy'

The only other mention of this bug/error/behavior I’ve seen so far is here (tf-github-issue) but it leads to nowhere unfortunately.

I am working in the official tensorflow docker image on tensorflow/tensorflow:2.11.1-gpu. So, at least I assume no dependency issues are present. Hopefully, somebody has seen this and can help me out, please!

Expected behavior is that the model loads as expected. As this is pretty common logic, I assume at this point that this an undocumented bug in tensorflow in relation to keras. I paid attention to using consistent namespaces such as import tensorflow as tf and use tensorflow functions as such tf.keras.models.load_model.