Fixed CUDA 11.8 workers

2023-12-20 07:50:14 +00:00
parent 210e513f96
commit c59038e902
3 changed files with 4 additions and 4 deletions
@@ -5,7 +5,7 @@ from vllm.model_executor.weight_utils import prepare_hf_model_weights
 if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=str)
-    parser.add_argument("--download_dir", type=str)
+    parser.add_argument("--download_dir", type=str, default=os.environ.get("MODEL_BASE_PATH"))

    args = parser.parse_args()
    if not args.model or not args.download_dir: