[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

bergson/collector/collector.py

@@ -333,9 +333,9 @@ def setup(self) -> None:
         """Initialize gradient storage dictionary."""
         self.mod_grads = {}
 
-        assert isinstance(self.model.device, torch.device), (
-            "Model device is not set correctly"
-        )
+        assert isinstance(
+            self.model.device, torch.device
+        ), "Model device is not set correctly"
 
         self.save_dtype = (
             torch.float32 if self.model.dtype == torch.float32 else torch.float16

bergson/gradients.py

@@ -512,9 +512,8 @@ def _process_grad(self, module: nn.Module, _, grad_out):
         G = grad_out[0]  # [N, S, O]
         I = module._inputs  # [N, S, I/q]
 
-
         name = assert_type(str, module._name)
-        if name== "h.1.mlp.c_fc":
+        if name == "h.1.mlp.c_fc":
             debugpy.breakpoint()
 
         # different way of checking for bias as above