Skip to content

Commit

Permalink
Fixed style issue.
Browse files Browse the repository at this point in the history
  • Loading branch information
HuFY-dev committed Mar 29, 2024
1 parent 18d7e8a commit 7cbbf82
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion sparse_autoencoder/optimizer/adam_with_reset.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,16 @@
"""
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 typing_extensions import TypeAlias

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]]
Expand Down

0 comments on commit 7cbbf82

Please sign in to comment.