Stoping and restarting training script starts from scratch

Prerequisites

Please answer the following questions for yourself before submitting an issue.

• I am using the latest TensorFlow Model Garden release and TensorFlow 2.
• I am reporting the issue to the correct repository. (Model Garden official or research directory)
• I checked to make sure that this issue has not already been filed.

1. The entire URL of the file you are using

https://github.com/tensorflow/models/tree/master/research/object_detection/model_main_tf2.py

2. Describe the bug

Contrary to TF1.x, in TF2.x when I stop training after a checkpoint, run evaluation and restart training the result is the model starts learning from scratch.

3. Steps to reproduce

1. Run training and produce a checkpoint.
2. Stop training and run evaluation.
3. Restart training.

Same thing happens when you skip step 2.

4. Expected behavior

I expect the training to continue from where it left of.

5. Additional context

Eficientnet D1 from https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2_detection_zoo.md

Default config and checkpoints.

6. System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 18.04.5 (tensorflow:2.3.0-gpu docker img)
• Mobile device name if the issue happens on a mobile device:
• TensorFlow installed from (source or binary): 2.3.0 from docker
• TensorFlow version (use command below): v2.3.0-rc2-23-gb36436b087 2.3.0
• Python version: Python 3.6.9
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version: from docker img
• GPU model and memory: from docker img

Checkpoint loading errors:

...
WARNING:tensorflow:Unresolved object in checkpoint: (root).model._box_predictor._base_tower_layers_for_heads.class_predictions_with_background.4.4.moving_variance
W0911 13:07:41.965147 140196924864320 util.py:150] Unresolved object in checkpoint: (root).model._box_predictor._base_tower_layers_for_heads.class_predictions_with_background.4.4.moving_variance
WARNING:tensorflow:Unresolved object in checkpoint: (root).model._box_predictor._base_tower_layers_for_heads.class_predictions_with_background.4.7.axis
W0911 13:07:41.965251 140196924864320 util.py:150] Unresolved object in checkpoint: (root).model._box_predictor._base_tower_layers_for_heads.class_predictions_with_background.4.7.axis
WARNING:tensorflow:Unresolved object in checkpoint: (root).model._box_predictor._base_tower_layers_for_heads.class_predictions_with_background.4.7.gamma
W0911 13:07:41.965382 140196924864320 util.py:150] Unresolved object in checkpoint: (root).model._box_predictor._base_tower_layers_for_heads.class_predictions_with_background.4.7.gamma
WARNING:tensorflow:Unresolved object in checkpoint: (root).model._box_predictor._base_tower_layers_for_heads.class_predictions_with_background.4.7.beta
W0911 13:07:41.965474 140196924864320 util.py:150] Unresolved object in checkpoint: (root).model._box_predictor._base_tower_layers_for_heads.class_predictions_with_background.4.7.beta
WARNING:tensorflow:Unresolved object in checkpoint: (root).model._box_predictor._base_tower_layers_for_heads.class_predictions_with_background.4.7.moving_mean
W0911 13:07:41.965547 140196924864320 util.py:150] Unresolved object in checkpoint: (root).model._box_predictor._base_tower_layers_for_heads.class_predictions_with_background.4.7.moving_mean
WARNING:tensorflow:Unresolved object in checkpoint: (root).model._box_predictor._base_tower_layers_for_heads.class_predictions_with_background.4.7.moving_variance
W0911 13:07:41.965606 140196924864320 util.py:150] Unresolved object in checkpoint: (root).model._box_predictor._base_tower_layers_for_heads.class_predictions_with_background.4.7.moving_variance
WARNING:tensorflow:A checkpoint was restored (e.g. tf.train.Checkpoint.restore or tf.keras.Model.load_weights) but not all checkpointed values were used. See above for specific issues. Use expect_partial() on the load status object, e.g. tf.train.Checkpoint.restore(...).expect_partial(), to silence these warnings, or use assert_consumed() to make the check explicit. See https://www.tensorflow.org/guide/checkpoint#loading_mechanics for details.
W0911 13:07:41.965777 140196924864320 util.py:158] A checkpoint was restored (e.g. tf.train.Checkpoint.restore or tf.keras.Model.load_weights) but not all checkpointed values were used. See above for specific issues. Use expect_partial() on the load status object, e.g. tf.train.Checkpoint.restore(...).expect_partial(), to silence these warnings, or use assert_consumed() to make the check explicit. See https://www.tensorflow.org/guide/checkpoint#loading_mechanics for details.
...