update: PARAMS, STATE

Author: kozistr

pytorch_optimizer/types.py

@@ -1,6 +1,10 @@
-from typing import Any, Callable, Dict, Optional, Tuple
+from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Union
+
+import torch
 
 CLOSURE = Optional[Callable[[], float]]
 LOSS = Optional[float]
 BETAS = Tuple[float, float]
 DEFAULT_PARAMETERS = Dict[str, Any]
+PARAMS = Union[Iterable[torch.Tensor], Iterable[Dict[str, Any]]]
+STATE = Dict[str, Any]