Unable to train CenterNet MobileNetV2 FPN with transfer learning

Prerequisites

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

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

1. The entire URL of the file you are using

http://download.tensorflow.org/models/object_detection/tf2/20210210/centernet_mobilenetv2fpn_512x512_coco17_od.tar.gz

2. Describe the bug

I am trying to train the centernet_mobilenetv2_fpn_od using the following tutorial:
https://tensorflow-object-detection-api-tutorial.readthedocs.io/en/latest/training.html#preparing-the-workspace
I used this tutorial previously to train ssd_mobilenet_v2 and it worked fine. I followed the same steps for centernet, but I recieve the following error:

tensorflow.python.framework.errors_impl.InvalidArgumentError: indices[0] = 0 is not in [0, 0)

3. Steps to reproduce

Follow this guide:
https://tensorflow-object-detection-api-tutorial.readthedocs.io/en/latest/training.html#preparing-the-workspace
but using the centernet_mobiletv2 model form the model zoo.

4. Expected behavior

CenterNet trains through transfer learning.

5. Additional context

Traceback (most recent call last):
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/context.py", line 2113, in execution_mode
    yield
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/data/ops/iterator_ops.py", line 733, in _next_internal
    output_shapes=self._flat_output_shapes)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/gen_dataset_ops.py", line 2579, in iterator_get_next
    _ops.raise_from_not_ok_status(e, name)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py", line 6862, in raise_from_not_ok_status
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: indices[0] = 0 is not in [0, 0)
	 [[{{node GatherV2_8}}]]
	 [[MultiDeviceIteratorGetNextFromShard]]
	 [[RemoteCall]] [Op:IteratorGetNext]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "model_main_tf2.py", line 113, in <module>
    tf.compat.v1.app.run()
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/platform/app.py", line 40, in run
    _run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
  File "/usr/local/lib/python3.7/dist-packages/absl/app.py", line 303, in run
    _run_main(main, args)
  File "/usr/local/lib/python3.7/dist-packages/absl/app.py", line 251, in _run_main
    sys.exit(main(argv))
  File "model_main_tf2.py", line 110, in main
    record_summaries=FLAGS.record_summaries)
  File "/content/drive/.shortcut-targets-by-id/1MNY3RVQFiMLhbbLGrZjkYHlMWIQeK7FY/Colab Notebooks/Tensorflow/workspace/training_demo/object_detection/model_lib_v2.py", line 597, in train_loop
    train_input, unpad_groundtruth_tensors)
  File "/content/drive/.shortcut-targets-by-id/1MNY3RVQFiMLhbbLGrZjkYHlMWIQeK7FY/Colab Notebooks/Tensorflow/workspace/training_demo/object_detection/model_lib_v2.py", line 392, in load_fine_tune_checkpoint
    _ensure_model_is_built(model, input_dataset, unpad_groundtruth_tensors)
  File "/content/drive/.shortcut-targets-by-id/1MNY3RVQFiMLhbbLGrZjkYHlMWIQeK7FY/Colab Notebooks/Tensorflow/workspace/training_demo/object_detection/model_lib_v2.py", line 155, in _ensure_model_is_built
    features, labels = iter(input_dataset).next()
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/input_lib.py", line 628, in next
    return self.__next__()
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/input_lib.py", line 632, in __next__
    return self.get_next()
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/input_lib.py", line 663, in get_next
    self._iterators[i].get_next_as_list_static_shapes(new_name))
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/input_lib.py", line 1619, in get_next_as_list_static_shapes
    return self._format_data_list_with_options(self._iterator.get_next())
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py", line 585, in get_next
    result.append(self._device_iterators[i].get_next())
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/data/ops/iterator_ops.py", line 800, in get_next
    return self._next_internal()
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/data/ops/iterator_ops.py", line 739, in _next_internal
    return structure.from_compatible_tensor_list(self._element_spec, ret)
  File "/usr/lib/python3.7/contextlib.py", line 130, in __exit__
    self.gen.throw(type, value, traceback)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/context.py", line 2116, in execution_mode
    executor_new.wait()
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/executor.py", line 69, in wait
    pywrap_tfe.TFE_ExecutorWaitForAllPendingNodes(self._handle)
tensorflow.python.framework.errors_impl.InvalidArgumentError: indices[0] = 0 is not in [0, 0)
	 [[{{node GatherV2_8}}]]
	 [[MultiDeviceIteratorGetNextFromShard]]
	 [[RemoteCall]]

6. System information

running on Google Collab. the following is installed before training:

• tf_slim
• lvis
• tf-models-official
• tensorflow-gpu

python -c "import tensorflow as tf; print(tf.version.GIT_VERSION, tf.version.VERSION)"
yeilds:

v2.4.1-0-g85c8b2a817f 2.4.1