pytorch/torch/optim/optimizer.pyi
Michael Kösel 4e0d098ace Fix optimizer type hint (#20648)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/20548
Pull Request resolved: https://github.com/pytorch/pytorch/pull/20648

Differential Revision: D15453935

Pulled By: ezyang

fbshipit-source-id: 8778e819c58fdc2620f123ec5b5fd568e23b7705
2019-05-22 11:27:40 -07:00

13 lines
477 B
Python

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