Unable to build trained object detection model restored using `checkpoint`

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. -> 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

2. Describe the bug

• I am using ssd_resnet50_v1_fpn_640x640_coco17_tpu-8 model
• After loading and modifying the config I have trained the model on kaggle using:

!python /kaggle/working/models/research/object_detection/model_main_tf2.py \
    --pipeline_config_path={pipeline_config_path} \
    --model_dir={model_dir} \
    --num_train_steps={num_steps} \
    --num_eval_steps={num_val_steps}

• Then I exported the model using:

!python /kaggle/working/models/research/object_detection/exporter_main_v2.py \
    --input_type=image_tensor \
    --pipeline_config_path=/kaggle/working/pipeline.config \
    --trained_checkpoint_dir=/kaggle/working/training --output_directory=/kaggle/working/exported_model

• To evaluate the model, I loaded the config -> model_config -> detection_model (using model_builder) -> restored checkpoint from above exported path.
• When I run a dummy image on the model I get the following error:

ValueError                                Traceback (most recent call last)
/tmp/ipykernel_41/2278393839.py in <module>
      1 tmp_img, tmp_shape = detection_model.preprocess(tf.zeros((1, 640, 640, 3)))
----> 2 tmp_predictions = detection_model.predict(tmp_img, tmp_shape)
      3 detections = detection_model.postprocess(tmp_predictions, tmp_shape)

/opt/conda/lib/python3.7/site-packages/object_detection/meta_architectures/ssd_meta_arch.py in predict(self, preprocessed_inputs, true_image_shapes)
    568       batchnorm_updates_collections = tf.GraphKeys.UPDATE_OPS
    569     if self._feature_extractor.is_keras_model:
--> 570       feature_maps = self._feature_extractor(preprocessed_inputs)
    571     else:
    572       with slim.arg_scope([slim.batch_norm],

/opt/conda/lib/python3.7/site-packages/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
   1035         with autocast_variable.enable_auto_cast_variables(
   1036             self._compute_dtype_object):
-> 1037           outputs = call_fn(inputs, *args, **kwargs)
   1038 
   1039         if self._activity_regularizer:

/opt/conda/lib/python3.7/site-packages/object_detection/meta_architectures/ssd_meta_arch.py in call(self, inputs, **kwargs)
    249   # method.
    250   def call(self, inputs, **kwargs):
--> 251     return self._extract_features(inputs)
    252 
    253 

/opt/conda/lib/python3.7/site-packages/object_detection/models/ssd_resnet_v1_fpn_keras_feature_extractor.py in _extract_features(self, preprocessed_inputs)
    223 
    224     image_features = self.classification_backbone(
--> 225         ops.pad_to_multiple(preprocessed_inputs, self._pad_to_multiple))
    226 
    227     feature_block_list = []

/opt/conda/lib/python3.7/site-packages/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
   1035         with autocast_variable.enable_auto_cast_variables(
   1036             self._compute_dtype_object):
-> 1037           outputs = call_fn(inputs, *args, **kwargs)
   1038 
   1039         if self._activity_regularizer:

/opt/conda/lib/python3.7/site-packages/keras/engine/functional.py in call(self, inputs, training, mask)
    413     """
    414     return self._run_internal_graph(
--> 415         inputs, training=training, mask=mask)
    416 
    417   def compute_output_shape(self, input_shape):

/opt/conda/lib/python3.7/site-packages/keras/engine/functional.py in _run_internal_graph(self, inputs, training, mask)
    548 
    549         args, kwargs = node.map_arguments(tensor_dict)
--> 550         outputs = node.layer(*args, **kwargs)
    551 
    552         # Update tensor_dict.

/opt/conda/lib/python3.7/site-packages/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
   1018         training=training_mode):
   1019 
-> 1020       input_spec.assert_input_compatibility(self.input_spec, inputs, self.name)
   1021       if eager:
   1022         call_fn = self.call

/opt/conda/lib/python3.7/site-packages/keras/engine/input_spec.py in assert_input_compatibility(input_spec, inputs, layer_name)
    216                          'expected ndim=' + str(spec.ndim) + ', found ndim=' +
    217                          str(ndim) + '. Full shape received: ' +
--> 218                          str(tuple(shape)))
    219     if spec.max_ndim is not None:
    220       ndim = x.shape.rank

ValueError: Input 0 of layer conv1_bn is incompatible with the layer: expected ndim=4, found ndim=1. Full shape received: (0,)

3. Steps to reproduce

Steps to reproduce are same as above

4. Expected behavior

It should have successfully build the model.

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): Kaggle
• Mobile device name if the issue happens on a mobile device: N.A
• TensorFlow installed from (source or binary): hopefully installed on conda from source
• TensorFlow version (use command below): 2.6.0
• Python version: 3.7.10
• Bazel version (if compiling from source):
• GCC/Compiler version (if compiling from source):
• CUDA/cuDNN version: 11.0
• GPU model and memory: NVIDIA Tesla (15781MiB / 16280MiB)

I wonder why the input received by the model lose their shape. Am I doing something wrong? Or is it a bug? Please do let me know.