Resolved import error for PyTorch >= 2.2

sparse_autoencoder/optimizer/adam_with_reset.py

@@ -2,16 +2,20 @@
 
 This reset method is useful when resampling dead neurons during training.
 """
-from collections.abc import Iterator
+from collections.abc import Iterable, Iterator
+from typing import Any
+from typing_extensions import TypeAlias
 
 from jaxtyping import Float, Int
 from torch import Tensor
 from torch.nn.parameter import Parameter
 from torch.optim import Adam
-from torch.optim.optimizer import params_t
 
 from sparse_autoencoder.tensor_types import Axis
 
+# params_t was renamed to ParamsT in PyTorch 2.2, which caused import errors
+# Copied from PyTorch 2.2 with modifications for better style
+ParamsT: TypeAlias = Iterable[Tensor] | Iterable[dict[str, Any]]
 
 class AdamWithReset(Adam):
     """Adam Optimizer with a reset method.
@@ -35,7 +39,7 @@ class AdamWithReset(Adam):
 
     def __init__(  # (extending existing implementation)
         self,
-        params: params_t,
+        params: ParamsT,
         lr: float | Float[Tensor, Axis.names(Axis.SINGLE_ITEM)] = 1e-3,
         betas: tuple[float, float] = (0.9, 0.999),
         eps: float = 1e-8,