Lint fix

kronfluence/computer/score_computer.py

@@ -8,8 +8,6 @@
 
 from kronfluence.arguments import FactorArguments, ScoreArguments
 from kronfluence.computer.computer import Computer
-from kronfluence.module.tracked_module import ModuleMode
-from kronfluence.module.utils import set_mode
 from kronfluence.score.pairwise import (
     compute_pairwise_scores_with_loaders,
     load_pairwise_scores,

kronfluence/module/conv2d.py

@@ -69,10 +69,12 @@ class TrackedConv2d(TrackedModule, module_type=nn.Conv2d):
 
     @property
     def weight(self) -> torch.Tensor:
+        """Returns the weight matrix."""
         return self.original_module.weight
 
     @property
     def bias(self) -> torch.Tensor:
+        """Returns the bias."""
         return self.original_module.bias
 
     def _get_flattened_activation(

kronfluence/module/linear.py

@@ -13,10 +13,12 @@ class TrackedLinear(TrackedModule, module_type=nn.Linear):
 
     @property
     def weight(self) -> torch.Tensor:
+        """Returns the weight matrix."""
         return self.original_module.weight
 
     @property
     def bias(self) -> torch.Tensor:
+        """Returns the bias."""
         return self.original_module.bias
 
     def _get_flattened_activation(

kronfluence/module/tracked_module.py

@@ -571,15 +571,14 @@ def _compute_low_rank_preconditioned_gradient(
                 torch.matmul(U_k, torch.diag_embed(S_k)).to(dtype=self.score_args.score_dtype).contiguous().clone(),
                 V_k.to(dtype=self.score_args.score_dtype),
             ]
-        else:
-            U, S, V = torch.svd_lowrank(
-                preconditioned_gradient.contiguous().to(dtype=self.score_args.query_gradient_svd_dtype),
-                q=rank,
-            )
-            return [
-                torch.matmul(U, torch.diag_embed(S)).to(dtype=self.score_args.score_dtype).contiguous().clone(),
-                V.transpose(1, 2).contiguous().to(dtype=self.score_args.score_dtype),
-            ]
+        U, S, V = torch.svd_lowrank(
+            preconditioned_gradient.contiguous().to(dtype=self.score_args.query_gradient_svd_dtype),
+            q=rank,
+        )
+        return [
+            torch.matmul(U, torch.diag_embed(S)).to(dtype=self.score_args.score_dtype).contiguous().clone(),
+            V.transpose(1, 2).contiguous().to(dtype=self.score_args.score_dtype),
+        ]
 
     def _compute_preconditioned_gradient(self, per_sample_gradient: torch.Tensor) -> None:
         """Computes the preconditioned per-sample-gradient.

tests/utils.py

@@ -59,7 +59,7 @@ def custom_scores_name(name: str) -> str:
 def prepare_model_and_analyzer(model: nn.Module, task: Task) -> Tuple[nn.Module, Analyzer]:
     model = prepare_model(model=model, task=task)
     analyzer = Analyzer(
-        analysis_name=f"pytest",
+        analysis_name="pytest",
         model=model,
         task=task,
         disable_model_save=True,