From e9490b425714e5261deae9136896dc4e2c353f0d Mon Sep 17 00:00:00 2001
From: gs-olive <113141689+gs-olive@users.noreply.github.com>
Date: Mon, 18 Dec 2023 11:47:42 -0800
Subject: [PATCH] fix: Docker builder with new Torch versions

- Version upgrade bump, narrow specification of versions
- Docker builder switch to GPU-provisioned machine
---
 .github/workflows/docker_builder.yml | 2 +-
 py/requirements.txt                  | 4 ++--
 pyproject.toml                       | 8 +++-----
 3 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/docker_builder.yml b/.github/workflows/docker_builder.yml
index c27757aae1..ebc3b4564a 100644
--- a/.github/workflows/docker_builder.yml
+++ b/.github/workflows/docker_builder.yml
@@ -16,7 +16,7 @@ concurrency:
 
 jobs:
   build:
-    runs-on: linux.2xlarge
+    runs-on: linux.4xlarge.nvidia.gpu
 
     # Define key environment variables
     # Container name is of the form torch_tensorrt:<branch_name>
diff --git a/py/requirements.txt b/py/requirements.txt
index 5883a4a878..cd52d32436 100644
--- a/py/requirements.txt
+++ b/py/requirements.txt
@@ -2,8 +2,8 @@ numpy
 packaging
 pybind11==2.6.2
 --extra-index-url https://download.pytorch.org/whl/nightly/cu121
-torch>=2.2.0.dev,<=2.3.0
-torchvision>=0.17.0.dev,<0.18.0
+torch>=2.3.0.dev,<2.4.0
+torchvision>=0.18.0.dev,<0.19.0
 --extra-index-url https://pypi.ngc.nvidia.com
 tensorrt==8.6.1
 pyyaml
diff --git a/pyproject.toml b/pyproject.toml
index 5de6c38fdf..c987ac1f40 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,8 +9,7 @@ requires = [
     "typing-extensions>=4.7.0",
     "future>=0.18.3",
     "tensorrt>=8.6,<8.7",
-    "torch >=2.2.0.dev,<=2.3.0",
-    #"torch==2.1.0.dev20230731",
+    "torch >=2.3.0.dev,<2.4.0",
     "pybind11==2.6.2",
     "numpy",
 ]
@@ -42,8 +41,7 @@ readme = {file = "py/README.md", content-type = "text/markdown"}
 requires-python = ">=3.8"
 keywords = ["pytorch", "torch", "tensorrt", "trt", "ai", "artificial intelligence", "ml", "machine learning", "dl", "deep learning", "compiler", "dynamo", "torchscript", "inference"]
 dependencies = [
-    "torch >=2.2.0.dev,<=2.3.0",
-    #"torch==2.1.0.dev20230731",
+    "torch >=2.3.0.dev,<2.4.0",
     "tensorrt>=8.6,<8.7",
     "packaging>=23",
     "numpy",
@@ -52,7 +50,7 @@ dependencies = [
 dynamic = ["version"]
 
 [project.optional-dependencies]
-torchvision = ["torchvision >=0.17.dev,<0.18.0"]
+torchvision = ["torchvision >=0.18.dev,<0.19.0"]
 
 [project.urls]
 Homepage = "https://pytorch.org/tensorrt"