try partial cache restore

jbloomAus · Apr 11, 2024 · 392f982 · 392f982
1 parent 89e1568
commit 392f982
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -55,6 +55,8 @@ jobs:
         with:
           path: .venv
           key: venv-${{ runner.os }}-${{ matrix.python-version }}-${{ hashFiles('**/pyproject.toml') }}
+          restore-keys: |
+            venv-${{ runner.os }}-${{ matrix.python-version }}-
       - name: Install dependencies
         if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true'
         run: poetry install --no-interaction

diff --git a/tests/unit/training/test_sparse_autoencoder.py b/tests/unit/training/test_sparse_autoencoder.py
@@ -200,7 +200,7 @@ def test_per_item_mse_loss_with_norm_matches_original_implementation() -> None:
         / (target_centered**2).sum(dim=-1, keepdim=True).sqrt()
     )
     sae_res = _per_item_mse_loss_with_target_norm(input, target)
-    assert torch.allclose(orig_impl_res, sae_res)
+    assert torch.allclose(orig_impl_res, sae_res, atol=1e-5)
 
 
 def test_SparseAutoencoder_remove_gradient_parallel_to_decoder_directions() -> None: