mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
595e344769
Summary: Forgot to mirror the `nn/ __init__.py` semantics in the new `nn` type stub. Pull Request resolved: https://github.com/pytorch/pytorch/pull/22411 Differential Revision: D16149798 Pulled By: ezyang fbshipit-source-id: 0ffa256fbdc5e5383a7b9c9c3ae61acd11de1dba
29 lines
696 B
Python
29 lines
696 B
Python
from typing import Any, TypeVar
|
|
from ..modules import Module
|
|
|
|
|
|
class WeightNorm:
|
|
name: str = ...
|
|
dim: int = ...
|
|
|
|
def __init__(self, name: str, dim: int) -> None: ...
|
|
|
|
# TODO Make return type more specific
|
|
def compute_weight(self, module: Module) -> Any: ...
|
|
|
|
@staticmethod
|
|
def apply(module: Module, name: str, dim: int) -> 'WeightNorm': ...
|
|
|
|
def remove(self, module: Module) -> None: ...
|
|
|
|
def __call__(self, module: Module, inputs: Any) -> None: ...
|
|
|
|
|
|
T_module = TypeVar('T_module', bound=Module)
|
|
|
|
|
|
def weight_norm(module: T_module, name: str = ..., dim: int = ...) -> T_module: ...
|
|
|
|
|
|
def remove_weight_norm(module: T_module, name: str = ...) -> T_module: ...
|