Home
last modified time | relevance | path

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

/external/pytorch/torch/distributed/
D_composable_state.py10 _module_state_mapping: Dict[nn.Module, _State] = {} variable
14 global _module_state_mapping
15 assert module not in _module_state_mapping, f"Inserting {module} more than once."
16 _module_state_mapping[module] = state
29 global _module_state_mapping
34 if module in _module_state_mapping:
35 return _module_state_mapping[module]