diff --git a/sdk/python/Dockerfile.conformance b/sdk/python/Dockerfile.conformance
index ad7a8d018e..35d4fc1838 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
+FROM continuumio/miniconda3
 
 WORKDIR /kubeflow/training-operator/sdk/python
 
@@ -25,6 +25,6 @@ COPY sdk/python/conformance/run.sh .
 RUN chmod +x run.sh
 
 RUN pip install pytest
-RUN python -m pip install -e .
+RUN conda install .
 
 ENTRYPOINT [ "./run.sh" ]
\ No newline at end of file
diff --git a/sdk/python/setup.py b/sdk/python/setup.py
index 0255d42137..9a7c650f8d 100644
--- a/sdk/python/setup.py
+++ b/sdk/python/setup.py
@@ -36,6 +36,9 @@
     "peft>=0.3.0",
     "datasets>=2.13.2",
     "huggingface_hub>=0.16.4",
+    "pytorch>=1.13.1",
+    "torchvision>=0.9.1",
+    "torchaudio>=0.8.1",
 ]
 
 setuptools.setup(