diff --git a/.github/workflows/cmake.yaml b/.github/workflows/cmake.yaml
index 5916a94..6c6f677 100644
--- a/.github/workflows/cmake.yaml
+++ b/.github/workflows/cmake.yaml
@@ -7,17 +7,19 @@ on:
     branches: [main]
 
 jobs:
-  tests:
+  unit-tests:
     runs-on: ubuntu-latest
+    container:
+      image: pytorch/pytorch:2.2.0
 
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
 
     - name: Install third-party libraries
       run: |
-        sudo apt-get update
-        sudo apt-get install -y libunwind-dev --no-install-recommends
-        sudo apt-get install -y libboost-test-dev libfmt-dev libtorch-dev --no-install-recommends
+        apt-get update
+        apt-get install -y libunwind-dev --no-install-recommends
+        apt-get install -y libboost-test-dev libfmt-dev libtorch-dev --no-install-recommends
 
     - name: Configure CMake
       run: cmake -B ${{github.workspace}}/build