diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml
index 4c8ce89..56131f5 100644
--- a/.github/workflows/unit-tests.yml
+++ b/.github/workflows/unit-tests.yml
@@ -14,6 +14,7 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
+        device: ["cpu", "cuda"]
 
     steps:
     - uses: actions/checkout@v2
@@ -21,15 +22,28 @@ jobs:
       uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
-    - name: Setup poetry and pytest
+    - name: Setup poetry
       run: |
         python -m pip install --upgrade pip
         pip install poetry
-        pip install pytest
-    - name: Install dependencies
+    - name: Install PyTorch
       run: |
         poetry env use python
-        poetry update
+        if [ "${{ matrix.device }}" == "cpu"]; then
+          if [ "${{ matrix.python-version }}" == "3.7"]; then
+            poetry run pip install torch==1.13.1+cpu torchvision==0.14.1+cpu torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cpu
+          else
+            poetry run pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+        elif [ "${{ matrix.device }}" == "cuda"]; then
+          if [ "${{ matrix.python-version }}" == "3.7"]; then
+            poetry run pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
+          else
+            poetry run pip install torch torchvision torchaudio
+        else
+          echo "Unsupported value for device"
+          exit 1
+    - name: Install DebyeCalculator dependencies
+        poetry install
     - name: Run pytest
       run: |
         poetry run pytest