// Set the location of call op to QuantizationUnitLoc if found.
  Location call_op_loc = location;
  for (Operation* op : cloning_ops) {
    std::optional<QuantizationUnitLoc::QuantizationUnit> unit =
        FindQuantizationUnitFromLoc(op->getLoc());
    if (unit.has_value()) {
      call_op_loc = QuantizationUnitLoc(builder.getContext(), unit.value());
    }
  }

      # Verify if quant_unit.pb file was created correctly.
      quant_unit_file_path = os.path.join(log_dir_path, folder, 'quant_unit.pb')

      quant_unit = (
          quant_opts_pb2.UnitWiseQuantizationSpec.QuantizationUnit.FromString(
              open(quant_unit_file_path, 'rb').read()
          )
      )

      self.assertEqual(quant_unit.node_name, 'Conv2D')