Allow installation of PyTorch with different backends

requirements.txt

@@ -1,6 +1,5 @@
 unsloth
 peft
-torch
 transformers
 datasets
 trl

setup.sh

@@ -9,6 +9,40 @@ echo ""
 echo "Creating main virtual environment..."
 python -m venv venv
 source venv/bin/activate
+
+# Select PyTorch version
+echo ""
+echo "Select PyTorch version:"
+echo "  1) CUDA (NVIDIA GPU, default)"
+echo "  2) ROCm (AMD GPU)"
+echo "  3) Custom version"
+echo ""
+read -p "Enter choice (1-3): " TORCH_CHOICE
+
+case $TORCH_CHOICE in
+    1)
+        echo "Installing PyTorch (CUDA)..."
+        pip install torch
+        ;;
+    2)
+        echo "Installing PyTorch (ROCm)..."
+        pip install torch --index-url https://download.pytorch.org/whl/nightly/rocm7.2
+        ;;
+    3)
+        echo "Checking for custom PyTorch..."
+        python -c "import torch" 2>/dev/null
+        if [ $? -ne 0 ]; then
+            echo "Error: torch is not available. Please install it before running this script."
+            exit 1
+        fi
+        echo "torch is available, continuing..."
+        ;;
+    *)
+        echo "Invalid choice. Installing PyTorch (CUDA) as default."
+        pip install torch
+        ;;
+esac
+
 pip install -r requirements.txt
 
 # Ask if fresh build or existing