diff --git a/custom-nb-image/Dockerfile b/custom-nb-image/Dockerfile
index 7da9b62ef..2fd9c02b5 100644
--- a/custom-nb-image/Dockerfile
+++ b/custom-nb-image/Dockerfile
@@ -26,5 +26,4 @@ RUN pip uninstall pickle5 -y
 RUN pip install codeflare-sdk==0.4.0 \
 				datasets==2.6.1 \
 				transformers==4.23.1 \
-				evaluate==0.3.0 \
-				git+https://github.com/project-codeflare/torchx.git@OCP
+				evaluate==0.3.0
diff --git a/pyproject.toml b/pyproject.toml
index e2ffd6227..061cd997d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -25,4 +25,4 @@ openshift-client = "1.0.18"
 rich = "^12.5"
 ray = {version = "2.1.0", extras = ["default"]}
 kubernetes = "26.1.0"
-torchx = {git = "https://github.com/project-codeflare/torchx", rev = "OCP"}
+codeflare-torchx = "0.5.0.dev5"
diff --git a/requirements.txt b/requirements.txt
index a17a43ad2..a805aacd8 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,4 +2,4 @@ openshift-client==1.0.18
 rich==12.5.1
 ray[default]==2.1.0
 kubernetes==26.1.0
-git+https://github.com/project-codeflare/torchx@OCP
+codeflare-torchx==0.5.0.dev5
diff --git a/src/codeflare_sdk.egg-info/PKG-INFO b/src/codeflare_sdk.egg-info/PKG-INFO
index 56985dd97..c4061c623 100644
--- a/src/codeflare_sdk.egg-info/PKG-INFO
+++ b/src/codeflare_sdk.egg-info/PKG-INFO
@@ -1,3 +1,4 @@
 Metadata-Version: 2.1
 Name: codeflare-sdk
 Version: 0.0.0
+License-File: LICENSE
diff --git a/src/codeflare_sdk.egg-info/SOURCES.txt b/src/codeflare_sdk.egg-info/SOURCES.txt
index 4a6c82bd6..96222b989 100644
--- a/src/codeflare_sdk.egg-info/SOURCES.txt
+++ b/src/codeflare_sdk.egg-info/SOURCES.txt
@@ -1,3 +1,4 @@
+LICENSE
 README.md
 pyproject.toml
 src/codeflare_sdk/__init__.py
@@ -6,9 +7,12 @@ src/codeflare_sdk.egg-info/SOURCES.txt
 src/codeflare_sdk.egg-info/dependency_links.txt
 src/codeflare_sdk.egg-info/top_level.txt
 src/codeflare_sdk/cluster/__init__.py
+src/codeflare_sdk/cluster/auth.py
 src/codeflare_sdk/cluster/cluster.py
 src/codeflare_sdk/cluster/config.py
 src/codeflare_sdk/cluster/model.py
+src/codeflare_sdk/job/__init__.py
+src/codeflare_sdk/job/jobs.py
 src/codeflare_sdk/utils/__init__.py
 src/codeflare_sdk/utils/generate_yaml.py
 src/codeflare_sdk/utils/pretty_print.py
\ No newline at end of file