def augment_graph(self):
        """
        Make all quantization_candidates op type nodes as part of the graph output.
        :return: augmented ONNX model
        """
        # for ORT version >= 1.16.0, we need directly use the model without clone
        # since clone_model_with_shape_infer is removed in ORT 1.16.0
        if is_ort_version_below_1_16():
            from onnxruntime.quantization.quant_utils import clone_model_with_shape_infer

            model = clone_model_with_shape_infer(self.model)
        else:
            model = self.model

        self.tensors_to_calibrate, value_infos = self.select_tensors_to_calibrate(model)
        for tensor in self.tensors_to_calibrate:
            if tensor not in self.model_original_outputs:
                model.graph.output.append(value_infos[tensor])

        onnx.save(
            model,
            self.augmented_model_path,
            save_as_external_data=self.use_external_data_format,
        )
        self.augment_model = model

        Make all quantization_candidates op type nodes as part of the graph output.
        :return: augmented ONNX model
        
olive/passes/onnx/vitis_ai/calibrate.py