use_dynamic_shape=true ran into inference failure of TensorRT when running tf serving on GPU A100

Description

I have a TF model which internally has some data processing layers,

class ScatterNdUpdater(tf.keras.layers.Layer):
      def __init__(self, **kwargs):
          super(ScatterNdUpdater, self).__init__(**kwargs)
      
      def call(self, tensor, indices_without_batch):
          **batch_size = tf.shape(indices_without_batch)[0]**
          **indices = tf.stack([tf.range(0, batch_size, dtype=tf.int32), tf.cast(indices_without_batch, tf.dtypes.int32)], axis=-1)**
          values = tf.ones((batch_size,))
          return tf.tensor_scatter_nd_update(tensor, indices, values)

where indices_without_batch is calculated by

class ReduceOneHotToFloat(tf.keras.layers.Layer):
    def __init__(self, feature_dimensionality, **kwargs):
          super(ReduceOneHotToFloat, self).__init__(**kwargs)
          self.feature_dimensionality = feature_dimensionality
          self.mask = tf.reshape(tf.range(0, feature_dimensionality, dtype=tf.float32), [1, -1])
    
    def call(self, one_hot_input):
          indices_without_batch = tf.math.argmax(one_hot_input, axis=-1)
          # This line is only needed for warmup, which a dummy batch of ones is passed, breaking one hot assumptions.
          indices_without_batch_clipped = tf.clip_by_value(indices_without_batch, 0, self.feature_dimensionality - 1)
          return indices_without_batch_clipped

TRT conversion is done by

trt_converter = trt.TrtGraphConverterV2(
input_saved_model_dir=self.input_model.model_path,
**use_dynamic_shape=True,**
dynamic_shape_profile_strategy="Range", # This will only be used if use_dynamic_shape is turned on
precision_mode=self.optimizer_config.precision_mode,
allow_build_at_runtime=self.optimizer_config.allow_build_at_runtime,
maximum_cached_engines=self.optimizer_config.maximum_cached_engines,
)
...
def input_fn():
for batch_size in (self.optimizer_config.min_batch_size, self.optimizer_config.max_batch_size):
logger.info(f"Optimizing for batch size: {batch_size}")
input_data = {
**generate_dummy_dense_tensors_from_model(batch_size, self.input_model.input_signature),
**generate_dummy_sparse_tensors_from_model(batch_size, self.input_model.input_signature),
}
input_data_converted = {convert_to_trt_feature_name(k): v for k, v in input_data.items()}
trt_model_input_features = trt_func._arg_keywords
vals = [input_data_converted[fea_name] for fea_name in trt_model_input_features]
yield vals

...
trt_converter.build(input_fn)
trt_converter.save(TMP_TF_TRT_MODEL_PATH)
...

Inference encountered 0.01% error rate of
INVALID_ARGUMENT: Shapes of all inputs must match: values[0].shape = [520] != values[1].shape = [540]

This happens only when use_dynamic_shapes = True, when it's false, there was no error.

Environment

TensorRT Version: 8.6.1

NVIDIA GPU: A100

NVIDIA Driver Version:

CUDA Version:

CUDNN Version:

Operating System:

Python Version (if applicable):

Tensorflow Version (if applicable):

PyTorch Version (if applicable):

Baremetal or Container (if so, version):

Relevant Files

Model link:

Steps To Reproduce

Commands or scripts:

Have you tried the latest release?:

Can this model run on other frameworks? For example run ONNX model with ONNXRuntime (polygraphy run <model.onnx> --onnxrt):

Start with the TrtGraphConverterV2 configuration using use_dynamic_shape=True and the provided input_fn, then reproduce the INVALID_ARGUMENT shape mismatch with the ScatterNdUpdater and ReduceOneHotToFloat layers. The issue is done when the dynamic-shape TensorRT build and inference no longer produce mismatched input shapes.