feat(test_train_sae): format with ruff

OpenMOSS · Oct 30, 2024 · c5f9dbf · c5f9dbf
1 parent 4cc893d
commit c5f9dbf
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/tests/intergration/test_attributor.py b/tests/intergration/test_attributor.py
@@ -9,6 +9,7 @@
 
 class TestModule(HookedRootModule):
     __test__ = False
+
     def __init__(self):
         super().__init__()
         self.W_1 = nn.Parameter(torch.tensor([[1.0, 2.0]]))

diff --git a/tests/intergration/test_train_sae.py b/tests/intergration/test_train_sae.py
@@ -1,9 +1,10 @@
-from transformer_lens import HookedTransformer, HookedTransformerConfig
 import torch
+from einops import rearrange
 from torch.optim import Adam
+from transformer_lens import HookedTransformer, HookedTransformerConfig
+
 from lm_saes.config import SAEConfig
 from lm_saes.sae import SparseAutoEncoder
-from einops import rearrange
 
 
 def test_train_sae():
@@ -41,9 +42,7 @@ def test_train_sae():
     ### Get activations ###
     tokens = torch.randint(0, 50, (batch_size, 10))
     with torch.no_grad():
-        _, cache = model.run_with_cache_until(
-            tokens, names_filter=hook_point, until=hook_point
-        )
+        _, cache = model.run_with_cache_until(tokens, names_filter=hook_point, until=hook_point)
         batch = {
             hook_point: rearrange(
                 cache[hook_point].to(dtype=dtype, device=device),