Chapter 12 loading model WARNING in colab

When running the [jupyter notebook](https://colab.research.google.com/github/ageron/handson-ml2/blob/master/12_custom_models_and_training_with_tensorflow.ipynb) in Colaboratory, tf is unable to load the saved optimizer state.

```
model.save("my_model_with_a_custom_loss.h5",)
model = keras.models.load_model("my_model_with_a_custom_loss.h5",
custom_objects={"huber_fn": huber_fn})
```
The message is
```
WARNING:tensorflow:Error in loading the saved optimizer state. As a result, your model is starting with a freshly initialized optimizer.
```
I wonder if anyone has encountered with the same issue, and any ideas? Tones of thanks!