From a66259972ae13bcf39867844dcb48b8e93de1a91 Mon Sep 17 00:00:00 2001
From: lanluo-nvidia <lanl@nvidia.com>
Date: Sun, 15 Sep 2024 11:01:28 -0700
Subject: [PATCH] release 2.5 branch cut (#3161)

---
 docker/dist-build.sh                         | 4 ++--
 py/requirements.txt                          | 6 +++---
 pyproject.toml                               | 8 ++++----
 tests/py/dynamo/models/test_models_export.py | 7 +++----
 tests/py/requirements.txt                    | 3 +--
 5 files changed, 13 insertions(+), 15 deletions(-)

diff --git a/docker/dist-build.sh b/docker/dist-build.sh
index 04d2e8b84a..19c3af8d50 100755
--- a/docker/dist-build.sh
+++ b/docker/dist-build.sh
@@ -5,9 +5,9 @@ set -x
 TOP_DIR=$(cd $(dirname $0); pwd)/..
 
 if [[ -z "${USE_CXX11}" ]]; then
-    BUILD_CMD="python -m pip wheel .  --extra-index-url https://download.pytorch.org/whl/nightly/cu124 -w dist"
+    BUILD_CMD="python -m pip wheel .  --extra-index-url https://download.pytorch.org/whl/test/cu124 -w dist"
 else
-    BUILD_CMD="python -m pip wheel . --config-setting="--build-option=--use-cxx11-abi" --extra-index-url https://download.pytorch.org/whl/nightly/cu124 -w dist"
+    BUILD_CMD="python -m pip wheel . --config-setting="--build-option=--use-cxx11-abi" --extra-index-url https://download.pytorch.org/whl/test/cu124 -w dist"
 fi
 
 # TensorRT restricts our pip version
diff --git a/py/requirements.txt b/py/requirements.txt
index 15d359fbc3..af3f8ef9ad 100644
--- a/py/requirements.txt
+++ b/py/requirements.txt
@@ -1,8 +1,8 @@
 numpy
 packaging
 pybind11==2.6.2
---extra-index-url https://download.pytorch.org/whl/nightly/cu124
-torch>=2.5.0.dev,<2.6.0
-torchvision>=0.20.0.dev,<0.21.0
+--extra-index-url https://download.pytorch.org/whl/test/cu124
+torch>=2.5.0,<2.6.0
+torchvision>=0.20.0,<0.21.0
 --extra-index-url https://pypi.ngc.nvidia.com
 pyyaml
diff --git a/pyproject.toml b/pyproject.toml
index f6230c8a74..55d7c3d506 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,7 +9,7 @@ requires = [
     "typing-extensions>=4.7.0",
     "future>=0.18.3",
     "tensorrt-cu12==10.3.0",
-    "torch >=2.5.0.dev,<2.6.0",
+    "torch>=2.5.0,<2.6.0",
     "pybind11==2.6.2",
     "numpy",
 ]
@@ -54,7 +54,7 @@ keywords = [
     "inference",
 ]
 dependencies = [
-    "torch >=2.5.0.dev,<2.6.0",
+    "torch>=2.5.0,<2.6.0",
     "tensorrt-cu12==10.3.0",
     "tensorrt-cu12-bindings==10.3.0",
     "tensorrt-cu12-libs==10.3.0",
@@ -65,8 +65,8 @@ dependencies = [
 dynamic = ["version"]
 
 [project.optional-dependencies]
-torchvision = ["torchvision >=0.20.dev,<0.21.0"]
-quantization = ["nvidia-modelopt[all]>=0.15.1"]
+torchvision = ["torchvision>=0.20.0,<0.21.0"]
+quantization = ["nvidia-modelopt[deploy,hf,torch]~=0.17.0"]
 monitoring-tools = ["rich >= 13.7.1"]
 jupyter = ["rich[jupyter] >= 13.7.1"]
 
diff --git a/tests/py/dynamo/models/test_models_export.py b/tests/py/dynamo/models/test_models_export.py
index 8976c7fbc0..d547a2e0bf 100644
--- a/tests/py/dynamo/models/test_models_export.py
+++ b/tests/py/dynamo/models/test_models_export.py
@@ -1,5 +1,5 @@
 # type: ignore
-import importlib
+import platform
 import unittest
 from importlib import metadata
 
@@ -250,9 +250,8 @@ def calibrate_loop(model):
 
 
 @unittest.skipIf(
-    not importlib.util.find_spec("modelopt")
-    or Version(metadata.version("nvidia-modelopt")) < Version("0.16.1"),
-    "modelopt 0.16.1 or later is required Int8 quantization is supported in modelopt since 0.16.1 or later",
+    modelopt.__version__ < "0.16.1",
+    "Int8 quantization is supported in modelopt since 0.16.1 or later",
 )
 @pytest.mark.unit
 def test_base_int8(ir):
diff --git a/tests/py/requirements.txt b/tests/py/requirements.txt
index 460cc01027..f791444c8a 100644
--- a/tests/py/requirements.txt
+++ b/tests/py/requirements.txt
@@ -9,6 +9,5 @@ pytest-xdist>=3.6.1
 pyyaml
 timm>=1.0.3
 transformers==4.40.2
-# TODO @lanlao-nvidia Renable when modelopt can be install properly to run the tests
-# "nvidia-modelopt[all]">=0.16.1,<0.17.0
+nvidia-modelopt[deploy,hf,torch]~=0.17.0
 --extra-index-url https://pypi.nvidia.com