),
        tags=tags,
        signature_keys=['serving_default'],
        op_set=quant_opts_pb2.XLA,
        # Gather op is opt-outed if the size is smaller than the threshold.
        min_num_elements_for_weights=1024 if expect_quantized_gather else 8192,
    )

    data_gen = self._create_data_generator(
        input_key='input_tensor',
        shape=[6],
        minval=0,
        maxval=10,