handle ampere and non ampere installs

install.sh

@@ -1,5 +1,20 @@
 #!/bin/bash
 
+echo "Please select your GPU generation:"
+echo "1. Ampere or newer (for newer GPUs)"
+echo "2. Non-Ampere (for older GPUs)"
+read -p "Enter your choice (1 or 2): " GPU_CHOICE
+
+if [[ "$GPU_CHOICE" == "1" ]]; then
+    echo "You have selected Ampere or newer. Proceeding with Ampere-specific configurations."
+    GPU_TYPE="ampere"
+elif [[ "$GPU_CHOICE" == "2" ]]; then
+    echo "You have selected Non-Ampere. Proceeding with configurations for older GPUs."
+    GPU_TYPE=""
+else
+    echo "Invalid selection. Please restart the script and choose either 1 or 2."
+fi
+
 # Function to get PyTorch version if installed
 get_pytorch_version () {
     if python -c "import torch" &> /dev/null; then
@@ -23,7 +38,11 @@ conda_install_packages () {
 # Function to install packages via Pip
 pip_install_packages () {
     pip install --upgrade --force-reinstall --no-cache-dir torch==${PYTORCH_CORE_VERSION}+${CUDA_TAG} triton --index-url https://download.pytorch.org/whl/${CUDA_TAG}
-    pip install "unsloth[${CUDA_TAG}] @ git+https://github.com/unslothai/unsloth.git"
+    if [[ "$PYTORCH_VERSION_TAG" == "210" ]]; then
+        pip install "unsloth[${CUDA_TAG}${GPU_TYPE:+-$GPU_TYPE}] @ git+https://github.com/unslothai/unsloth.git"
+    else
+        pip install "unsloth[${CUDA_TAG}${GPU_TYPE:+-$GPU_TYPE}-$PYTORCH_VERSION_TAG] @ git+https://github.com/unslothai/unsloth.git"
+    fi
 }
 
 # Check if conda is installed
@@ -69,6 +88,7 @@ else
 
         # Extract PyTorch version (ignoring any suffix)
         PYTORCH_CORE_VERSION=$(echo $PYTORCH_VERSION | cut -d'+' -f1)
+        PYTORCH_VERSION_TAG="torch${PYTORCH_CORE_VERSION//./}"
 
         pip_install_packages