Add ruff fixes

sharanry · Oct 27, 2024 · 8ce0d7f · 8ce0d7f
1 parent cac38fe
commit 8ce0d7f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/probe_lens/probes.py b/probe_lens/probes.py
@@ -1,11 +1,11 @@
-from tqdm import tqdm
 import torch
 import torch.nn as nn
+from tqdm.autonotebook import tqdm
 
 
 class LinearProbe(nn.Module):
     def __init__(self, input_dim, output_dim=1, device="cpu"):
-        super(LinearProbe, self).__init__()
+        super().__init__()
         self.linear = nn.Linear(input_dim, output_dim, device=device)
 
     def forward(self, x):

diff --git a/tests/test_probes.py b/tests/test_probes.py
@@ -1,7 +1,8 @@
-from probe_lens.probes import LinearProbe
 import torch
 import torch.nn as nn
 
+from probe_lens.probes import LinearProbe
+
 
 def test_linear_probe():
     input_dim = 10