def _convert_model_on_device(
        self,
        model: Union[HfModelHandler, PyTorchModelHandler],
        config: type[BasePassConfig],
        output_model_path: str,
        device: str,
        torch_dtype: Optional[torch.dtype] = None,
    ) -> ONNXModelHandler:
        """Convert an HfModelHandler or PyTorchModelHandler to an ONNXModelHandler."""
        # prepare the model for conversion
        if isinstance(model, HfModelHandler):
            # optimum export config needs the loaded model to get io_config so we create a new model handler
            # which will be used to load the model and get the io_config
            model = self._prepare_hf_model(model, device, torch_dtype)

        # load the model
        pytorch_model = model.load_model(cache_model=False)
        pytorch_model.eval()

        # get dummy inputs
        dummy_inputs = self._get_dummy_inputs(model, config)
        io_config = model.io_config

        converted_onnx_model = OnnxConversion._export_pytorch_model(
            pytorch_model, dummy_inputs, io_config, config, device, torch_dtype, tempfile.tempdir
        )

        model_attributes = deepcopy(model.model_attributes or {})

        # add split information if present
        split_assignments = model_attributes.get("split_assignments")
        if split_assignments:
            split_assignment_str = ";".join([f"{k}={v}" for k, v in split_assignments.items()])
            onnx.helper.set_model_props(converted_onnx_model, {"split_assignments": split_assignment_str})

        # save the model to the output path and return the model
        output_model_path = resolve_onnx_path(output_model_path)
        output_model = model_proto_to_olive_model(converted_onnx_model, output_model_path, config)
        output_model.model_attributes = deepcopy(model.model_attributes or {})
        return output_model
Convert an HfModelHandler or PyTorchModelHandler to an ONNXModelHandler.
olive/passes/onnx/conversion.py