Lines Matching defs:Module

44     root_module: nn.Module,
45     fn: Callable[[nn.Module], Optional[nn.Module]],
58         module: nn.Module,
60         parent_module: Optional[nn.Module],
85     root_module: nn.Module,
86     target_module_to_kwargs: Dict[nn.Module, Dict[str, Any]],
95     def fn(module: nn.Module) -> Optional[nn.Module]:
107     root_module: nn.Module,
108     module_classes: Iterable[Type[nn.Module]],
109     ignored_modules: Set[nn.Module],
111     target_module_to_kwargs: Dict[nn.Module, Dict[str, Any]],
144         root_module: nn.Module,
145         ignored_modules: Set[nn.Module],
156     module: nn.Module,
159     module_classes: Set[Type[nn.Module]],
193     def __init__(self, module_classes: Iterable[Type[nn.Module]]):
200         root_module: nn.Module,
201         ignored_modules: Set[nn.Module],
249     def __init__(self, lambda_fn: Callable[[nn.Module], Union[bool, Dict[str, Any]]]):
254         root_module: nn.Module,
255         ignored_modules: Set[nn.Module],
280     module: nn.Module, recurse: bool, nonwrapped_numel: int, lambda_fn: Callable
308     module: nn.Module,
311     transformer_layer_cls: Set[Type[nn.Module]],
323     module: nn.Module, module_classes: Sequence[type], recurse: bool, *args, **kwargs
335     module: nn.Module,
351     module: nn.Module,
356     force_leaf_modules: Optional[Set[Type[nn.Module]]] = None,
357     exclude_wrap_modules: Optional[Set[Type[nn.Module]]] = None,
445 def wrap(module: nn.Module, **wrap_overrides: Any) -> nn.Module:
482 def _wrap(module: nn.Module, wrapper_cls: Callable, **kwargs) -> nn.Module:
496     module: nn.Module,
499     ignored_modules: Set[nn.Module],