diff --git a/build_scripts/inference/serve b/build_scripts/inference/serve
index a57df749..fb64fe82 100755
Binary files a/build_scripts/inference/serve and b/build_scripts/inference/serve differ
diff --git a/build_scripts/install_comfy.sh b/build_scripts/install_comfy.sh
index faac187b..6e87c0c1 100755
--- a/build_scripts/install_comfy.sh
+++ b/build_scripts/install_comfy.sh
@@ -59,7 +59,7 @@ if [ "$ON_SAGEMAKER" == "true" ]; then
   pip install onnxruntime
 
 #  pip install torch==2.0.1 torchvision==0.15.2 --extra-index-url https://download.pytorch.org/whl/cu118
-  pip install torch==2.4.0 torchvision==0.19.0 --extra-index-url https://download.pytorch.org/whl/cu118
+  pip install torch==2.4.0 torchvision==0.19.0 --extra-index-url https://download.pytorch.org/whl/cu121
   pip install https://github.com/openai/CLIP/archive/d50d76daa670286dd6cacf3bcd80b5e4823fc8e1.zip
   pip install https://github.com/mlfoundations/open_clip/archive/bb6e834e9c70d9c27d0dc3ecedeebeaeb1ffad6b.zip
   pip install open-clip-torch==2.20.0
@@ -82,7 +82,7 @@ else
   /venv/bin/python3 -s -m pip install onnxruntime
 
 #  /venv/bin/python3 -m pip install torch==2.0.1 torchvision==0.15.2 --extra-index-url https://download.pytorch.org/whl/cu118
-  /venv/bin/python3 -m pip install torch==2.4.0 torchvision==0.19.0 --extra-index-url https://download.pytorch.org/whl/cu118
+  /venv/bin/python3 -m pip install torch==2.4.0 torchvision==0.19.0 --extra-index-url https://download.pytorch.org/whl/cu121
   /venv/bin/python3 -m pip install https://github.com/openai/CLIP/archive/d50d76daa670286dd6cacf3bcd80b5e4823fc8e1.zip
   /venv/bin/python3 -m pip install https://github.com/mlfoundations/open_clip/archive/bb6e834e9c70d9c27d0dc3ecedeebeaeb1ffad6b.zip
   /venv/bin/python3 -m pip install open-clip-torch==2.20.0