WeightSharedConvolutionalBoxPredictor box_prediction_head not in self._prediction_heads

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/blob/master/research/object_detection/predictors/convolutional_keras_box_predictor.py

2. Describe the bug

In Line 286 the class_prediction_head is put in self._prediction_heads, but not box_prediction_head.
Is that a bug that it's missing, or intended? And if intended, then why is it missing and why is it in self._prediction_heads in ConvolutionalBoxPredictor in Line 121?

3. Steps to reproduce

Steps to reproduce the behavior.

4. Expected behavior

A clear and concise description of what you expected to happen.

5. Additional context

Include any logs that would be helpful to diagnose the problem.

6. System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
• Mobile device name if the issue happens on a mobile device:
• TensorFlow installed from (source or binary):
• TensorFlow version (use command below):
• Python version:
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version:
• GPU model and memory: