feature: PARAMS, PARAM_GROUP, PARAM_GROUPS

Author: kozistr

pytorch_optimizer/types.py

@@ -1,10 +1,12 @@
-from typing import Any, Callable, Dict, Iterable, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Iterable, List, 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]]]
+PARAMS = Union[Iterable[Dict[str, Any]], Iterable[torch.Tensor]]
+PARAM_GROUP = Dict
+PARAM_GROUPS = List[PARAM_GROUP]
 STATE = Dict[str, Any]