Cannot restore a checkpoint on a pruned model without 'Unresolved object in checkpoint: (root).optimizer.*'

Prior to filing: check that this should be a bug instead of a feature request. Everything supported, including the compatible versions of TensorFlow, is listed in the overview page of each technique. For example, the overview page of quantization-aware training is here. An issue for anything not supported should be a feature request.

Describe the bug
Cannot load checkpoints into a saved model and restore the optimiser. Model weights are restored but it fails to restore optimiser. I get warnings:

WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.iter
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.iter
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.beta_1
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.beta_1
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.beta_2
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.beta_2
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.decay
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.decay
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.learning_rate
WARNING:tensorflow:Unresolved object in checkpoint: (root).optimizer.learning_rate

Here is a Colab that reproduces the problem. It trains MNIST, prunes it with checkpoints then tries to reload each checkpoint after training on the saved h5.

https://colab.research.google.com/drive/1Oljrqs0IHwDfEBfz8i-2aIdsfvH5Nkfd?usp=sharing

System information

TensorFlow version (installed from source or binary):
2.3 binary presumably

TensorFlow Model Optimization version (installed from source or binary):
0.5 binary presumably

Python version:

Describe the expected behavior
Restore the optimiser without warnings

Describe the current behavior
Warns

Code to reproduce the issue
Provide a reproducible code that is the bare minimum necessary to generate the
problem.

https://colab.research.google.com/drive/1Oljrqs0IHwDfEBfz8i-2aIdsfvH5Nkfd?usp=sharing

Start by running the linked Colab reproduction with the reported TensorFlow 2.3 and TensorFlow Model Optimization 0.5 versions. Trace the pruning checkpoint save and reload steps, focusing on why model weights restore while optimizer objects remain unresolved. Done means the reproduced checkpoint restores the optimizer without the listed warnings.