Searched defs:offload_to_cpu (Results 1 – 4 of 4) sorted by relevance
| /aosp_15_r20/external/pytorch/test/distributed/fsdp/ |
| H A D | test_hsdp_dtensor_state_dict.py | 110 def test_dtensor_sharded_tensor_state_dict_identical(self, offload_to_cpu): argument
179 def test_dtensor_sharded_optim_load_state_dict(self, offload_to_cpu): argument
233 def test_dtensor_sharded_model_load_state_dict(self, offload_to_cpu): argument
|
| H A D | test_fsdp_dtensor_state_dict.py | 115 self, offload_to_cpu, is_even_sharded_model argument
197 self, offload_to_cpu, is_even_sharded_model argument
255 self, offload_to_cpu, is_even_sharded_model argument
|
| /aosp_15_r20/external/pytorch/torch/distributed/fsdp/ |
| H A D | api.py | 288 offload_to_cpu: bool = False variable in StateDictConfig
367 offload_to_cpu: bool = True variable in OptimStateDictConfig
384 offload_to_cpu: bool = False variable in LocalOptimStateDictConfig
|
| /aosp_15_r20/external/pytorch/torch/testing/_internal/distributed/ |
| H A D | common_state_dict.py | 23 def _compare_tensor(self, orig_tensor, dist_tensor, offload_to_cpu=False): argument
38 offload_to_cpu=False, argument
|