diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 28fd0023..dc363baa 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -14,7 +14,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [openstack22]
+        os: [openstack22, tpu]
         python: ["3.12"]
     timeout-minutes: 500
     steps:
@@ -22,16 +22,17 @@ jobs:
         run: |
             if [ "$(expr substr $(uname -s) 1 5)" == "Linux" ]; then
               sudo apt-get update
-              sudo apt-get install -y git gcc g++ #python3 python3-dev python3-pip
-              #sudo python3 -m pip install --upgrade lit
-              #- run: python3 -m pip install --user numpy
+              sudo apt-get install -y git gcc g++ 
             fi
       - uses: actions/checkout@v3
         with:
           submodules: recursive
       
-      - run: sudo rm -rf ~/.cache/bazel ~/.cache/bazelisk
-      
+      - run: |
+            if [ "${{matrix.os}}" == "openstack22" ]; then
+              sudo rm -rf ~/.cache/bazel ~/.cache/bazelisk
+            fi
+
       - uses: bazel-contrib/setup-bazel@0.8.5
         name: Set up Bazel
         with:
@@ -46,6 +47,11 @@ jobs:
       - name: test
         run: |
           HERMETIC_PYTHON_VERSION=${{ matrix.python }} bazel test --test_output=errors ...
+          HERMETIC_PYTHON_VERSION=${{ matrix.python }} bazel test --cache_test_results=no //test:bench_vs_xla
+          HERMETIC_PYTHON_VERSION=${{ matrix.python }} bazel test --cache_test_results=no //test:llama
+          cat bazel-out/*/testlogs/test/llama/test.log
+          cp bazel-out/*/testlogs/test/llama/test.log llama.log
+
           
       - name: Build Wheel
         run: |
@@ -57,4 +63,4 @@ jobs:
         uses: actions/upload-artifact@v3
         with:
           name: Wheel-${{ matrix.os }}
-          path: "*.whl"
+          path: "*.{whl,log}"
diff --git a/test/BUILD b/test/BUILD
index 6558078e..5b43a375 100644
--- a/test/BUILD
+++ b/test/BUILD
@@ -108,6 +108,7 @@ py_test(
     ],
     imports = ["."],
     deps = TEST_DEPS,
+    timeout='long'
 )
 
 py_test(