diff --git a/.github/actions/smoke-test/check-image-size.sh b/.github/actions/smoke-test/check-image-size.sh
index e1cad85d2..7a48ff69b 100755
--- a/.github/actions/smoke-test/check-image-size.sh
+++ b/.github/actions/smoke-test/check-image-size.sh
@@ -1,4 +1,7 @@
 #!/bin/bash
+
+set -e
+
 # Function to handle errors
 handle_error() {
     local exit_code=$?
@@ -8,7 +11,6 @@ handle_error() {
     exit $exit_code
 }
 trap 'handle_error $LINENO ${BASH_COMMAND%% *}' ERR
-echo "This is line $LINENO"
 
 convert_gb_to_bytes() {
     local gb="$1"
@@ -52,6 +54,7 @@ check_image_size() {
     if [ -n $image_size  ] && [ $image_size -gt $threshold ]; then
         echo -e "\nImage size exceeds the threshold of $THRESHOLD_IN_GB gb"
         echo -e "\n❌ Image size check failed."
+        exit 1;
     else
         echo -e "\n✅  Passed!"
     fi
diff --git a/.github/workflows/smoke-universal.yaml b/.github/workflows/smoke-universal.yaml
index 108853972..62775c0ba 100644
--- a/.github/workflows/smoke-universal.yaml
+++ b/.github/workflows/smoke-universal.yaml
@@ -26,4 +26,4 @@ jobs:
       uses: ./.github/actions/smoke-test
       with:
         image: universal
-        threshold: 14
+        threshold: 10
diff --git a/src/universal/.devcontainer/local-features/machine-learning-packages/devcontainer-feature.json b/src/universal/.devcontainer/local-features/machine-learning-packages/devcontainer-feature.json
index 8687627d8..6f8bff36c 100644
--- a/src/universal/.devcontainer/local-features/machine-learning-packages/devcontainer-feature.json
+++ b/src/universal/.devcontainer/local-features/machine-learning-packages/devcontainer-feature.json
@@ -4,4 +4,4 @@
 	"installsAfter": [
 		"ghcr.io/devcontainers/features/python"
 	]
-}
+}
\ No newline at end of file
diff --git a/src/universal/.devcontainer/local-features/machine-learning-packages/install.sh b/src/universal/.devcontainer/local-features/machine-learning-packages/install.sh
index a86e30b78..c16153a2a 100644
--- a/src/universal/.devcontainer/local-features/machine-learning-packages/install.sh
+++ b/src/universal/.devcontainer/local-features/machine-learning-packages/install.sh
@@ -31,10 +31,10 @@ export DEBIAN_FRONTEND=noninteractive
 
 install_python_package() {
     PACKAGE=${1:-""}
-
+    OPTIONS=${2:-""}
     sudo_if /usr/local/python/current/bin/python -m pip uninstall --yes $PACKAGE
     echo "Installing $PACKAGE..."
-    sudo_if /usr/local/python/current/bin/python -m pip install --user --upgrade --no-cache-dir $PACKAGE
+    sudo_if /usr/local/python/current/bin/python -m pip install --user --upgrade --no-cache-dir $PACKAGE $OPTIONS
 }
 
 if [[ "$(python --version)" != "" ]] && [[ "$(pip --version)" != "" ]]; then
@@ -44,11 +44,11 @@ if [[ "$(python --version)" != "" ]] && [[ "$(pip --version)" != "" ]]; then
     install_python_package "matplotlib"
     install_python_package "seaborn"
     install_python_package "scikit-learn"
-    install_python_package "torch"
+    install_python_package "torch" "-f https://download.pytorch.org/whl/cpu/torch_stable.html"
     install_python_package "requests"
     install_python_package "plotly"
 else
     "(*) Error: Need to install python and pip."
 fi
 
-echo "Done!"
+echo "Done!"
\ No newline at end of file