diff --git a/sdk/python/Dockerfile.conformance b/sdk/python/Dockerfile.conformance
index 4cdc6f3bd6..ad7a8d018e 100644
--- a/sdk/python/Dockerfile.conformance
+++ b/sdk/python/Dockerfile.conformance
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 # Dockerfile for building the source code of conformance tests
-FROM python:3.10-alpine
+FROM python:3.10
 
 WORKDIR /kubeflow/training-operator/sdk/python
 
diff --git a/sdk/python/setup.py b/sdk/python/setup.py
index f3e854fa83..3ea2865c39 100644
--- a/sdk/python/setup.py
+++ b/sdk/python/setup.py
@@ -75,4 +75,5 @@
     install_requires=REQUIRES,
     tests_require=TESTS_REQUIRES,
     extras_require={"test": TESTS_REQUIRES},
+    dependency_links="https://download.pytorch.org/whl/cu118",
 )