for (const auto& arg : quantization_driver.GetArgs()) {
    const QuantState& state = quantization_driver.GetArgQuantState(arg);
    EXPECT_TRUE(isa<quant::QuantizedType>(state.params));
  }
  for (const auto& result : quantization_driver.GetResultStates()) {
    Operation* op = result.first.first;
    const int res_index = result.first.second;
    const QuantState state =
        quantization_driver.GetResultQuantState(op, res_index);