tweak torch version check

mllam · Feb 5, 2025 · bd84926 · bd84926
1 parent 9d17b40
commit bd84926
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/.github/workflows/install-and-test.yml b/.github/workflows/install-and-test.yml
@@ -58,10 +58,19 @@ jobs:
         run: |
           ${{ matrix.package_manager == 'pdm' && 'pdm install --group :all' || 'python -m pip install ".[dev]"' }}
 
-      - name: Print and check torch version
+      - name: Print torch version
         run: |
           ${{ matrix.package_manager == 'pdm' && 'pdm run' || '' }} python -c "import torch; print(torch.__version__)"
-          ${{ matrix.package_manager == 'pdm' && 'pdm run' || '' }} python -c "import torch; assert ('+cpu' in torch.__version__) == (matrix.device == 'cpu')"
+
+      - name: Check torch version and variant is cpu
+        if: matrix.device == 'cpu'
+        run: |
+          ${{ matrix.package_manager == 'pdm' && 'pdm run' || '' }} python -c "import torch; assert torch.__version__.endswith('+cpu')"
+
+      - name: Check torch version and variant is gpu
+        if: matrix.device == 'gpu'
+        run: |
+          ${{ matrix.package_manager == 'pdm' && 'pdm run' || '' }} python -c "import torch; assert not torch.__version__.endswith('+cpu')"
 
       - name: Load cache data
         uses: actions/cache/restore@v4