Searched defs:StateDictType (Results 1 – 4 of 4) sorted by relevance
| /aosp_15_r20/external/pytorch/test/distributed/fsdp/ |
| H A D | test_fsdp_optim_state.py | 523 state_dict_type: StateDictType,
550 state_dict_type: StateDictType,
1142 state_dict_type: StateDictType,
1175 state_dict_type: StateDictType,
1258 state_dict_type: StateDictType,
1279 state_dict_type: StateDictType,
1532 def test_save_load_without_0th_param_state(self, state_dict_type: StateDictType):
|
| /aosp_15_r20/external/pytorch/torch/distributed/fsdp/ |
| H A D | fully_sharded_data_parallel.py | 645 state_dict_type: StateDictType,
801 state_dict_type: StateDictType,
|
| H A D | api.py | 244 class StateDictType(Enum): class
|
| /aosp_15_r20/external/pytorch/test/distributed/checkpoint/fsdp/ |
| H A D | test_fsdp_dsd.py | 123 def _test_save_with_fsdp1_and_load_with_fsdp2(self, state_dict_type: StateDictType):
|