@test_util.run_in_graph_and_eager_modes
  def test_qat_gather_and_conv_model(
      self,
  ):
    input_type = dtypes.int32
    model = self._create_simple_gather_and_conv_model(
        input_type,
        filter_shape=(2, 3, 3, 1024),
        is_qat_model=True,
    )

    saved_model_save.save(model, self._input_saved_model_path)

        sess,
        tags,
        signature_def_map={signature_key: sig_def},
        main_op=init_op,
        assets_collection=assets_collection,
    )
    v1_builder.save()

  def _create_simple_gather_and_conv_model(
      self,
      input_type: dtypes.DType,
      filter_shape: Sequence[int],
      is_qat_model: bool = False,
  ) -> module.Module:
    class SimpleGatherAndConvModel(module.Module):