test

Author: casteryh

.github/workflows/unit_test.yaml

@@ -6,7 +6,7 @@ on:
 
 jobs:
   unit_tests:
-    runs-on: ubuntu-latest
+    runs-on: 4-core-ubuntu-gpu-t4
     timeout-minutes: 15
     strategy:
       matrix:
@@ -23,8 +23,12 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Update pip
         run: python -m pip install --upgrade pip
+      - name: Check CUDA version
+        run: |
+          nvcc --version
+          nvidia-smi
       - name: Install pytorch
-        run: python -m pip install torch==2.9.0.dev20250826 --extra-index-url https://download.pytorch.org/whl/nightly/cpu
+        run: python -m pip install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu126
       - name: Install monarch
         run: python -m pip install monarch-no-torch==0.1.0.dev20250826 --find-links assets/ci
       - name: Install torchstore