diff --git a/optimum/exporters/onnx/__main__.py b/optimum/exporters/onnx/__main__.py index 1e36af06ade..7f943512de5 100644 --- a/optimum/exporters/onnx/__main__.py +++ b/optimum/exporters/onnx/__main__.py @@ -224,9 +224,9 @@ def main_export( original_task = task task = TasksManager.map_from_synonym(task) - framework = TasksManager.determine_framework(model_name_or_path, subfolder=subfolder, framework=framework) + framework = TasksManager.determine_framework(model_name_or_path, subfolder=subfolder, framework=framework, token=token) library_name = TasksManager.infer_library_from_model( - model_name_or_path, subfolder=subfolder, library_name=library_name + model_name_or_path, subfolder=subfolder, library_name=library_name, token=token ) torch_dtype = None diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py index c0221f7bf69..d6c06751974 100644 --- a/optimum/exporters/tasks.py +++ b/optimum/exporters/tasks.py @@ -1457,6 +1457,7 @@ def determine_framework( subfolder: str = "", framework: Optional[str] = None, cache_dir: str = HUGGINGFACE_HUB_CACHE, + token: str = None, ) -> str: """ Determines the framework to use for the export. @@ -1484,7 +1485,7 @@ def determine_framework( if framework is not None: return framework - all_files, request_exception = TasksManager.get_model_files(model_name_or_path, subfolder, cache_dir) + all_files, request_exception = TasksManager.get_model_files(model_name_or_path, subfolder, cache_dir, token=token) pt_weight_name = Path(WEIGHTS_NAME).stem pt_weight_extension = Path(WEIGHTS_NAME).suffix