Update inference.py

inference.py

@@ -9,15 +9,15 @@ from typing import List, Dict
 try:
     import vllm
 except ImportError:
-    # Check CUDA version and install the correct vllm version
-    cuda_version = torch.version.cuda
-    if cuda_version == "11.8":
-        vllm_version = "v0.6.1.post1"
-        pip_cmd = f"pip install https://github.com/vllm-project/vllm/releases/download/{vllm_version}/vllm-{vllm_version}+cu118-cp310-cp310-manylinux1_x86_64.whl --extra-index-url https://download.pytorch.org/whl/cu118"
-    else:
-        raise RuntimeError(f"Unsupported CUDA version: {cuda_version}")
-
-    subprocess.check_call([sys.executable, "-m", "pip", "install", pip_cmd])
+    # Install vllm with CUDA 11.8 support
+    vllm_version = "v0.6.1.post1"
+    pip_cmd = [
+        sys.executable, 
+        "-m", "pip", "install", 
+        f"https://github.com/vllm-project/vllm/releases/download/{vllm_version}/vllm-{vllm_version}+cu118-cp310-cp310-manylinux1_x86_64.whl",
+        "--extra-index-url", "https://download.pytorch.org/whl/cu118"
+    ]
+    subprocess.check_call(pip_cmd)
 # Import the necessary modules after installation
 from vllm import LLM, SamplingParams
 from vllm.utils import random_uuid