Fix PyTorch version compatibility and update ROCm Dockerfile image

- Change torch to stable 2.4.0, add torchvision for compatibility - Update Dockerfile to use specific ROCm PyTorch image version
2026-05-13 02:15:48 +03:00
parent 67da25d849
commit 221ecbe69d
2 changed files with 5 additions and 4 deletions
--- a/4
+++ b/4
@@ -1,5 +1,5 @@
-# Use ROCm PyTorch base image
-FROM rocm/pytorch:latest
+# Use ROCm PyTorch base image with compatible PyTorch
+FROM rocm/pytorch:2.4.0-rocm6.0-cxx11-ubuntu22.04

 # Set environment variables for ROCm and Python
 ENV PYTHONDONTWRITEBYTECODE=1