make type hints compatible with python 3.8

MadryLab · Nov 2, 2023 · 3c0696a · 3c0696a
1 parent 102edbf
commit 3c0696a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trak/gradient_computers.py b/trak/gradient_computers.py
@@ -197,7 +197,7 @@ def __init__(
         grad_dim: int,
         dtype: torch.dtype,
         device: torch.device,
-        grad_wrt: Optional[list[str]] = None,
+        grad_wrt: Optional[Iterable[str]] = None,
     ) -> None:
         super().__init__(model, task, grad_dim, dtype, device)
         self.load_model_params(model)