RuntimeError: __iter__() is only supported inside of tf.function or when eager execution is enabled.

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

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

2. Describe the bug

I can't load the dataset as it requires eager_execution, which I thought should already initiated in tf2

3. Steps to reproduce

I follow the guidance in here
Run model_main_tf2.py with ssd_mobilenet_v2 checkpoint from the model zoo

4. Expected behavior

> Traceback (most recent call last):
>   File "model_main_tf2.py", line 115, in <module>
>     tf.compat.v1.app.run()
>   File "/home/dtlam26/.local/lib/python3.7/site-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 112, in main
>     record_summaries=FLAGS.record_summaries)
>   File "/usr/local/lib/python3.7/dist-packages/object_detection/model_lib_v2.py", line 643, in train_loop
>     train_input_iter = iter(train_input)
>   File "/home/dtlam26/.local/lib/python3.7/site-packages/tensorflow/python/distribute/input_lib.py", line 1315, in __iter__
>     raise RuntimeError("__iter__() is only supported inside of tf.function "
> RuntimeError: __iter__() is only supported inside of tf.function or when eager execution is enabled.
> 

5. Additional context

Nothing else

6. System information

• OS Platform and Distribution: Linux Ubuntu 18.04
• TensorFlow installed from (source or binary): source
• TensorFlow version (use command below): 2.4
• Python version: 3.7
• CUDA/cuDNN version: 11.0/8.0.5