Home
last modified time | relevance | path

Searched refs:_insert_module_state (Results 1 – 3 of 3) sorted by relevance

/external/pytorch/torch/distributed/_composable/
Dfully_shard.py9 from torch.distributed._composable_state import _get_module_state, _insert_module_state
124 _insert_module_state(module, state)
130 _insert_module_state(submodule, state)
/external/pytorch/torch/distributed/
D_composable_state.py13 def _insert_module_state(module: nn.Module, state: _State) -> None: function
/external/pytorch/torch/distributed/_composable/fsdp/
D_fsdp_state.py25 _insert_module_state,
93 _insert_module_state(module, self)