Home
last modified time | relevance | path

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

/aosp_15_r20/external/pytorch/torch/distributed/fsdp/
H A Dfully_sharded_data_parallel.py120 class OptimStateKeyType(Enum): class
1689 optim_state_key_type: OptimStateKeyType, argument
1736 OptimStateKeyType.PARAM_NAME,
1737 OptimStateKeyType.PARAM_ID,
1750 optim_state_key_type == OptimStateKeyType.PARAM_NAME
1753 optim_state_key_type == OptimStateKeyType.PARAM_ID
1759 if optim_state_key_type == OptimStateKeyType.PARAM_NAME: # ID -> name
1782 elif optim_state_key_type == OptimStateKeyType.PARAM_ID: # name -> ID
H A D__init__.py12 OptimStateKeyType,
/aosp_15_r20/external/pytorch/test/distributed/fsdp/
H A Dtest_fsdp_optim_state.py23 OptimStateKeyType,
1312 OptimStateKeyType.PARAM_ID,
1319 OptimStateKeyType.PARAM_ID,
1380 OptimStateKeyType.PARAM_NAME,
1387 OptimStateKeyType.PARAM_NAME,
1515 OptimStateKeyType.PARAM_ID,
1525 OptimStateKeyType.PARAM_NAME,