Searched defs:fsdp_model (Results 1 – 3 of 3) sorted by relevance
| /aosp_15_r20/external/pytorch/torch/testing/_internal/ |
| H A D | common_fsdp.py | 1131 def _check_cpu_offload(self, fsdp_model, cpu_offload): argument
1134 def _check_backward_prefetch(self, fsdp_model, backward_prefetch): argument
1137 def _check_forward_prefetch(self, fsdp_model, forward_prefetch): argument
|
| /aosp_15_r20/external/pytorch/test/distributed/fsdp/ |
| H A D | test_fsdp_mixed_precision.py | 288 def _validate_no_mp_shard(self, fsdp_model): argument
298 def _validate_mp_shard_freed(self, fsdp_model): argument
|
| H A D | test_fsdp_comm.py | 83 def _run_iter(self, fsdp_model, batch, use_no_sync: bool): argument
|