From dee77232feb5c7e41216e5586da3ec4407638846 Mon Sep 17 00:00:00 2001 From: aarush gupta Date: Wed, 9 Oct 2024 13:03:16 -0700 Subject: [PATCH] fix type annotations (#1941) [skip ci] --- src/axolotl/monkeypatch/relora.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/axolotl/monkeypatch/relora.py b/src/axolotl/monkeypatch/relora.py index e4352cbe3..9d246cb17 100644 --- a/src/axolotl/monkeypatch/relora.py +++ b/src/axolotl/monkeypatch/relora.py @@ -44,8 +44,8 @@ def magnitude_pruning_(tensor, prune_ratio): def reset_optimizer( optimizer: torch.optim.Optimizer, *, - reset_params: list[str], # where str is the key to a torch.nn.Parameter - optimizer_state_keys: list[str], + reset_params: List[str], # where str is the key to a torch.nn.Parameter + optimizer_state_keys: List[str], prune_ratio: float = 0.9, ): pruning_fn = partial(magnitude_pruning_, prune_ratio=prune_ratio)