pytorch-mlir

Dockerfile

@@ -19,20 +19,17 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
     apt-get -y update && apt-get install -y mlir-18-tools llvm-18
 
-# Install PyTorch
-RUN --mount=type=cache,target=/root/.cache pip install --break-system-packages torch
+# Install pytorch-mlir
 
-# Install PyTorch-MLIR
-WORKDIR /opt
-RUN git clone https://github.com/llvm/torch-mlir.git
-WORKDIR /opt/torch-mlir
-RUN python3 -m pip install -e . --user
+RUN pip install --break-system-packages --pre torch-mlir torchvision \
+  --extra-index-url https://download.pytorch.org/whl/nightly/cpu \
+  -f https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels
 
 # Add PyTorch-MLIR tools to PATH
 ENV PATH="/root/.local/bin:${PATH}"
 
 # Verify installations
-RUN python3 -c "import torch; print(torch.__version__)" && \
+RUN python -c "import torch; print(torch.__version__)" && \
     mlir-opt --version && \
     llvm-as --version