Merge remote-tracking branch 'origin/develop' into feature/ekfac_notebook_new_framework

Author: Xuzzo

src/pydvl/influence/base_influence_function_model.py

@@ -37,7 +37,8 @@ def __init__(self):
 
 
 class NotImplementedLayerRepresentationException(ValueError):
-    def __init__(self, message: str):
+    def __init__(self, module_id: str):
+        message = f"Only Linear layers are supported, but found module {module_id} requiring grad."
         super().__init__(message)
 
 

src/pydvl/influence/torch/influence_function_model.py

@@ -957,9 +957,7 @@ def _init_layer_kfac_blocks(
             forward_x_layer = torch.zeros((sA, sA), device=module.weight.device)
             grad_y_layer = torch.zeros((sG, sG), device=module.weight.device)
         else:
-            raise NotImplementedLayerRepresentationException(
-                f"Only Linear layers are supported, but found module {module} requiring grad."
-            )
+            raise NotImplementedLayerRepresentationException(module_id=str(module))
         return forward_x_layer, grad_y_layer
 
     @staticmethod
@@ -993,9 +991,7 @@ def grad_hook(m, m_grad, m_out):
                 grad_y[m_name] += torch.mm(m_out.t(), m_out)
 
         else:
-            raise NotImplementedLayerRepresentationException(
-                f"Only Linear layers are supported, but found module {module} requiring grad."
-            )
+            raise NotImplementedLayerRepresentationException(module_id=str(module))
         return input_hook, grad_hook
 
     def _get_kfac_blocks(
@@ -1076,9 +1072,7 @@ def _init_layer_diag(module: torch.nn.Module) -> torch.Tensor:
             sA = module.in_features + int(with_bias)
             layer_diag = torch.zeros((sA * sG), device=module.weight.device)
         else:
-            raise NotImplementedLayerRepresentationException(
-                f"Only Linear layers are supported, but found module {module} requiring grad."
-            )
+            raise NotImplementedLayerRepresentationException(module_id=str(module))
         return layer_diag
 
     def _get_layer_diag_hooks(
@@ -1116,9 +1110,7 @@ def grad_hook(m, m_grad, m_out):
                 ).view(-1)
 
         else:
-            raise NotImplementedLayerRepresentationException(
-                f"Only Linear layers are supported, but found module {module} requiring grad."
-            )
+            raise NotImplementedLayerRepresentationException(module_id=str(module))
         return input_hook, grad_hook
 
     def _update_diag(