diff --git a/.github/workflows/build_pr_documentation.yml b/.github/workflows/build_pr_documentation.yml
index 1e7bd0a3..15d1bf81 100644
--- a/.github/workflows/build_pr_documentation.yml
+++ b/.github/workflows/build_pr_documentation.yml
@@ -31,6 +31,8 @@ jobs:
           sudo mv ./doc-build/optimum.amd ./doc-build/optimum-amd
           echo ${{ env.COMMIT_SHA }} > ./commit_sha
           echo ${{ env.PR_NUMBER }} > ./pr_number
+        env:
+          COMMIT_SHA_SUBPACKAGE: ${{ github.event.pull_request.head.sha }}
 
       - uses: actions/upload-artifact@v4
         with:
diff --git a/optimum/amd/brevitas/quantizer.py b/optimum/amd/brevitas/quantizer.py
index d06e143b..f24393f0 100644
--- a/optimum/amd/brevitas/quantizer.py
+++ b/optimum/amd/brevitas/quantizer.py
@@ -13,7 +13,7 @@
 from brevitas_examples.llm.llm_quant.equalize import apply_weight_equalization
 from tqdm import tqdm
 
-from optimum.exporters import TasksManager
+from optimum.exporters.tasks import TasksManager
 from optimum.quantization_base import OptimumQuantizer
 from transformers.utils.fx import symbolic_trace
 
diff --git a/optimum/amd/ryzenai/modeling.py b/optimum/amd/ryzenai/modeling.py
index 74186cef..a7f743ae 100644
--- a/optimum/amd/ryzenai/modeling.py
+++ b/optimum/amd/ryzenai/modeling.py
@@ -17,8 +17,8 @@
 from onnx import shape_inference
 from onnx.tools import update_model_dims
 
-from optimum.exporters import TasksManager
 from optimum.exporters.onnx import main_export
+from optimum.exporters.tasks import TasksManager
 from optimum.modeling_base import FROM_PRETRAINED_START_DOCSTRING, OptimizedModel
 from optimum.onnx.utils import _get_external_data_paths
 from optimum.utils.save_utils import maybe_load_preprocessors
diff --git a/optimum/amd/ryzenai/pipelines/__init__.py b/optimum/amd/ryzenai/pipelines/__init__.py
index 7b9c0051..2ba1755a 100644
--- a/optimum/amd/ryzenai/pipelines/__init__.py
+++ b/optimum/amd/ryzenai/pipelines/__init__.py
@@ -4,7 +4,7 @@
 
 from typing import TYPE_CHECKING, Any, Optional, Union
 
-from optimum.exporters import TasksManager
+from optimum.exporters.tasks import TasksManager
 from transformers import ImageClassificationPipeline, Pipeline, PretrainedConfig
 from transformers import pipeline as transformers_pipeline
 from transformers.image_processing_utils import BaseImageProcessor
diff --git a/setup.py b/setup.py
index 4eedceb7..f9ad971a 100644
--- a/setup.py
+++ b/setup.py
@@ -14,7 +14,7 @@
     assert False, "Error: Could not open '%s' due %s\n" % (filepath, error)
 
 # ORT 1.16 is not compatible: https://github.com/Xilinx/Vitis-AI/issues/1343
-INSTALL_REQUIRE = ["optimum", "transformers>=4.38", "onnx", "onnxruntime-extensions"]
+INSTALL_REQUIRE = ["optimum<2.0.0", "transformers>=4.38", "onnx", "onnxruntime-extensions"]
 
 # TODO: unpin pytest once https://github.com/huggingface/transformers/pull/29154 is merged & released
 TESTS_REQUIRE = [
diff --git a/tests/brevitas/test_onnx_export.py b/tests/brevitas/test_onnx_export.py
index 836d21b8..7f9a65a6 100644
--- a/tests/brevitas/test_onnx_export.py
+++ b/tests/brevitas/test_onnx_export.py
@@ -15,12 +15,12 @@
 from testing_utils import SUPPORTED_MODELS_TINY, VALIDATE_EXPORT_ON_SHAPES, get_quantized_model
 
 from optimum.amd.brevitas.export import find_and_insert_matmulinteger
-from optimum.exporters import TasksManager
 from optimum.exporters.onnx import (
     export_models,
     validate_models_outputs,
 )
 from optimum.exporters.onnx.utils import _get_submodels_and_onnx_configs
+from optimum.exporters.tasks import TasksManager
 from optimum.utils import DEFAULT_DUMMY_SHAPES
 from optimum.utils.testing_utils import grid_parameters
 from transformers.modeling_utils import get_parameter_dtype