fix import renaming on python 2.2.0

ai-safety-foundation · Feb 9, 2024 · 9c7a61a · 9c7a61a
1 parent f005cc1
commit 9c7a61a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/sparse_autoencoder/optimizer/adam_with_reset.py b/sparse_autoencoder/optimizer/adam_with_reset.py
@@ -8,7 +8,8 @@
 from torch import Tensor
 from torch.nn.parameter import Parameter
 from torch.optim import Adam
-from torch.optim.optimizer import params_t
+try: from torch.optim.optimizer import params_t
+except ImportError: from torch.optim.optimizer import ParamsT as params_t
 
 from sparse_autoencoder.tensor_types import Axis