optimizer.pyi 689 Bytes
Newer Older
1
	# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
2

3
from typing import Any, List, Dict, Iterable, Union, Callable, Optional
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
4
5
from .. import Tensor

6
_params_t = Union[Iterable[Tensor], Iterable[Dict]]
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
7
8

class Optimizer(object):
9
10
    param_groups: List[Dict]
    state: Dict
11
    def __init__(self, params: _params_t, defaults: Optional[Dict]=None, lr: Optional[float]=None) -> None: ...
12
13
    def state_dict(self) -> Dict: ...
    def load_state_dict(self, state_dict: Dict) -> None: ...
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
14
15
    def zero_grad(self) -> None: ...
    def step(self, closure: Optional[Callable[[], float]]=...) -> Optional[float]: ...
16
    def add_param_group(self, param_group: Dict) -> None: ...