Home
last modified time | relevance | path

Searched defs:CPUOffload (Results 1 – 13 of 13) sorted by relevance

/aosp_15_r20/external/pytorch/test/distributed/fsdp/
H A Dtest_fsdp_core.py78 def _get_cuda_init_modes(self, cpu_offload: CPUOffload) -> List[CUDAInitMode]:
92 def _get_subtest_config(self, cpu_offload: CPUOffload) -> Dict[str, List[Any]]:
110 cpu_offload: CPUOffload,
126 cpu_offload: CPUOffload,
149 cpu_offload: CPUOffload,
165 cpu_offload: CPUOffload,
181 cpu_offload: CPUOffload,
203 cpu_offload: CPUOffload,
229 cpu_offload: CPUOffload,
246 cpu_offload: CPUOffload,
H A Dtest_fsdp_sharded_grad_scaler.py170 cpu_offload: CPUOffload,
210 cpu_offload: CPUOffload = CPUOffload(offload_params=False),
254 cpu_offload: CPUOffload,
H A Dtest_fsdp_use_orig_params.py124 cpu_offload: CPUOffload,
342 cpu_offload: CPUOffload,
565 cpu_offload: CPUOffload,
635 cpu_offload: CPUOffload,
690 cpu_offload: CPUOffload,
H A Dtest_fsdp_checkpoint.py140 cpu_offload: CPUOffload,
206 cpu_offload: CPUOffload,
H A Dtest_fsdp_pure_fp16.py59 def _test_pure_fp16_training(self, cpu_offload: CPUOffload):
H A Dtest_fsdp_grad_acc.py94 cpu_offload: CPUOffload,
H A Dtest_fsdp_unshard_params.py121 cpu_offload: CPUOffload,
H A Dtest_wrap.py315 cpu_offload: CPUOffload,
/aosp_15_r20/external/pytorch/torch/distributed/_composable/
H A Dfully_shard.py55 cpu_offload: Optional[CPUOffload] = None,
/aosp_15_r20/external/pytorch/torch/distributed/fsdp/
H A Dapi.py230 class CPUOffload: class
H A D_init_utils.py433 cpu_offload: Optional[CPUOffload],
H A Dfully_sharded_data_parallel.py416 cpu_offload: Optional[CPUOffload] = None,
/aosp_15_r20/external/pytorch/torch/testing/_internal/
H A Dcommon_fsdp.py1203 fsdp_cpu_offload: Optional[CPUOffload] = None,
1292 cpu_offload: CPUOffload = CPUOffload(),