Searched refs:_init_model (Results 1 – 4 of 4) sorted by relevance
| /external/pytorch/torch/distributed/checkpoint/examples/ |
| D | stateful_example.py | 59 def _init_model(device, world_size): function
81 model, optim = _init_model(device, world_size)
90 model, optim = _init_model(device, world_size)
|
| D | async_checkpointing_example.py | 51 def _init_model(rank, world_size): function
89 model, optim = _init_model(rank, world_size)
|
| /external/pytorch/test/distributed/fsdp/ |
| D | test_fsdp_comm_hooks.py | 174 def _init_model(self, core, sharding_strategy, mixed_precision=None): member in TestCommunicationHooks
205 fsdp_model_with_hook = self._init_model(
255 fsdp_model_with_hook = self._init_model(
315 fsdp_model_with_hook = self._init_model(
342 fsdp_with_hook = self._init_model(
349 fsdp_with_mp = self._init_model(
|
| D | test_fsdp_comm.py | 55 def _init_model( member in TestCommunication
229 fsdp_model = self._init_model(nested_model, sharding_strategy, device)
|