fix broken cuda and rocm images

docker/cuda/Dockerfile

@@ -28,14 +28,18 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 
 # Install PyTorch
 ARG TORCH_CUDA=cu124
-ARG TORCH_VERSION=stable
+ARG TORCH_RELEASE_TYPE=stable
+ARG TORCH_VERSION
 
-RUN if [ "${TORCH_VERSION}" = "stable" ]; then \
-    pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
-elif [ "${TORCH_VERSION}" = "nighly" ]; then \
+RUN if [ -n "${TORCH_VERSION}" ]; then \
+    # Install specific torch version if TORCH_VERSION is provided
+    pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
+elif [ "${TORCH_RELEASE_TYPE}" = "stable" ]; then \
+    pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/${TORCH_CUDA} ; \
+elif [ "${TORCH_RELEASE_TYPE}" = "nightly" ]; then \
     pip install --no-cache-dir --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
 else \
-    pip install --no-cache-dir torch==${TORCH_VERSION} torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_CUDA} ; \
+    echo "Error: Invalid PYTORCH_RELEASE_TYPE. Must be 'stable', 'nightly', or specify a TORCH_VERSION." && exit 1 ; \
 fi
 
 # Install quantization libraries from source

scripts/install_quantization_libs.py

@@ -1,29 +1,11 @@
 import argparse
 import os
-import re
 import subprocess
 import sys
 
 EXTERNAL_REPOS_DIR = "external_repos"
 
 
-def process_setup_file(setup_file_path):
-    with open(setup_file_path, "r") as file:
-        setup_content = file.read()
-
-    # Use a regular expression to remove any line containing "torch=="
-    setup_content = re.sub(r'"torch==[^\"]+",', "", setup_content)
-
-    # Set IS_CPU_ONLY to False
-    setup_content = setup_content.replace(
-        "IS_CPU_ONLY = not torch.backends.mps.is_available() and not torch.cuda.is_available()", "IS_CPU_ONLY = False"
-    )
-
-    # Write the modified content back to setup.py
-    with open(setup_file_path, "w") as file:
-        file.write(setup_content)
-
-
 def clone_or_pull_repo(repo_url, repo_location_path):
     """Clone the repo if it doesn't exist; otherwise, pull the latest changes."""
     if os.path.exists(repo_location_path):
@@ -46,8 +28,6 @@ def install_autoawq_from_source():
     kernels_repo_path = os.path.join(EXTERNAL_REPOS_DIR, autoawq_kernels_repo_name)
 
     clone_or_pull_repo(f"https://github.com/casper-hansen/{autoawq_kernels_repo_name}", kernels_repo_path)
-    kernels_setup_file_path = os.path.join(kernels_repo_path, "setup.py")
-    process_setup_file(kernels_setup_file_path)
     subprocess.run(
         f"cd {kernels_repo_path} && {sys.executable} -m pip install .",
         shell=True,
@@ -56,8 +36,6 @@ def install_autoawq_from_source():
     )
 
     clone_or_pull_repo(f"https://github.com/casper-hansen/{autoawq_repo_name}", autoawq_repo_path)
-    autoawq_setup_file_path = os.path.join(autoawq_repo_path, "setup.py")
-    process_setup_file(autoawq_setup_file_path)
     subprocess.run(
         f"cd {autoawq_repo_path} && {sys.executable} -m pip install .",
         shell=True,
@@ -75,8 +53,6 @@ def install_autogptq_from_source():
 
     clone_or_pull_repo("https://github.com/PanQiWei/AutoGPTQ.git", autogptq_repo_path)
     subprocess.run("pip install numpy gekko pandas", shell=True, check=True, env=os.environ)
-    autogptq_setup_file_path = os.path.join(autogptq_repo_path, "setup.py")
-    process_setup_file(autogptq_setup_file_path)
     subprocess.run(
         f"cd {autogptq_repo_path} && {sys.executable} -m pip install .",
         shell=True,