Fixed import & style

ai-safety-foundation · Mar 29, 2024 · 1e7e3fc · 1e7e3fc
1 parent 659d60c
commit 1e7e3fc
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/sparse_autoencoder/optimizer/adam_with_reset.py b/sparse_autoencoder/optimizer/adam_with_reset.py
@@ -12,10 +12,11 @@
 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.1
-from typing import Union, Iterable, Dict, Any
+# Copied from PyTorch 2.1 with modifications for better style
+from collections.abc import Iterable
+from typing import Any
 from typing_extensions import TypeAlias
-params_t: TypeAlias = Union[Iterable[Tensor], Iterable[Dict[str, Any]]]
+params_t: TypeAlias = Iterable[Tensor] | Iterable[dict[str, Any]]
 
 class AdamWithReset(Adam):
     """Adam Optimizer with a reset method.