pytorch/torch/optim/optimizer.pyi
James Reed 57b54dfec5 Fix Optimizer.zero_grad type annotation (#76998)
`Optimizer.zero_grad()` defines the `set_to_none` argument as `bool`, not `Optional[bool]`

Fixes #ISSUE_NUMBER

Pull Request resolved: https://github.com/pytorch/pytorch/pull/76998
Approved by: https://github.com/albanD
2022-05-11 00:05:26 +00:00

19 lines
651 B
Python

from typing import Iterable, Union, Callable, Optional, List
from .. import Tensor
_params_t = Union[Iterable[Tensor], Iterable[dict]]
class Optimizer:
defaults: dict
state: dict
param_groups: List[dict]
def __init__(self, params: _params_t, default: dict) -> None: ...
def __setstate__(self, state: dict) -> None: ...
def state_dict(self) -> dict: ...
def load_state_dict(self, state_dict: dict) -> None: ...
def zero_grad(self, set_to_none: bool=...) -> None: ...
def step(self, closure: Optional[Callable[[], float]]=...) -> Optional[float]: ...
def add_param_group(self, param_group: dict) -> None: ...