Fix/435 broken imports

src/pydvl/influence/general.py

@@ -104,8 +104,9 @@ def test_grads() -> Generator[TensorType, None, None]:
             )  # type:ignore
 
     try:
-        # if provided input_data implements __len__, pre-allocate the result tensor to reduce memory consumption
-        resulting_shape = (len(test_data), model.num_params)  # type:ignore
+        # in case input_data is a torch DataLoader created from a Dataset,
+        # we can pre-allocate the result tensor to reduce memory consumption
+        resulting_shape = (len(test_data.dataset), model.num_params)  # type:ignore
         rhs = cat_gen(
             test_grads(), resulting_shape, model  # type:ignore
         )  # type:ignore
@@ -174,8 +175,9 @@ def train_grads() -> Generator[TensorType, None, None]:
             )  # type:ignore
 
     try:
-        # if provided input_data implements __len__, pre-allocate the result tensor to reduce memory consumption
-        resulting_shape = (len(input_data), model.num_params)  # type:ignore
+        # in case input_data is a torch DataLoader created from a Dataset,
+        # we can pre-allocate the result tensor to reduce memory consumption
+        resulting_shape = (len(input_data.dataset), model.num_params)  # type:ignore
         train_grad_tensor = cat_gen(
             train_grads(), resulting_shape, model  # type:ignore
         )  # type:ignore

src/pydvl/influence/torch/functional.py

@@ -91,7 +91,7 @@ def batch_hvp_gen(
 
     for inputs, targets in iter(data_loader):
         batch_loss = batch_loss_function(model, loss, inputs, targets)
-        model_params = dict(model.named_parameters())
+        model_params = {k: p for k, p in model.named_parameters() if p.requires_grad}
 
         def batch_hvp(vec: torch.Tensor):
             return flatten_tensors_to_vector(
@@ -166,9 +166,7 @@ def batch_loss_function(
     """
 
     def batch_loss(params: Dict[str, torch.Tensor]):
-        outputs = functional_call(
-            model, params, (to_model_device(x, model),), strict=True
-        )
+        outputs = functional_call(model, params, (to_model_device(x, model),))
         return loss(outputs, y)
 
     return batch_loss
@@ -209,7 +207,9 @@ def get_hvp_function(
     """
 
     params = {
-        k: p if track_gradients else p.detach() for k, p in model.named_parameters()
+        k: p if track_gradients else p.detach()
+        for k, p in model.named_parameters()
+        if p.requires_grad
     }
 
     def hvp_function(vec: torch.Tensor) -> torch.Tensor:

src/pydvl/influence/torch/torch_differentiable.py

@@ -149,7 +149,6 @@ def model_func(param):
                     param,
                 ),
                 (x.to(self.device),),
-                strict=True,
             )
             return self.loss(outputs, y.to(self.device))