Searched defs:fsdp_kwargs (Results 1 – 5 of 5) sorted by relevance
| /external/pytorch/test/distributed/fsdp/ |
| D | test_fsdp_freezing_weights.py | 41 fsdp_kwargs, argument
58 def fsdp_wrap(self, fsdp_kwargs): argument
74 fsdp_kwargs, argument
92 def fsdp_wrap(self, fsdp_kwargs): argument
126 fsdp_kwargs, argument
|
| D | test_fsdp_state_dict.py | 190 self, *fsdp_args, wrap=True, checkpoint_wrap=False, **fsdp_kwargs
209 def _get_simple_model(self, *fsdp_args, checkpoint_wrap=False, **fsdp_kwargs):
217 self, *fsdp_args, wrap=True, checkpoint_wrap=False, **fsdp_kwargs
250 def _get_non_fsdp_root_module(self, *fsdp_args, wrap=True, **fsdp_kwargs):
|
| D | test_fsdp_mixed_precision.py | 247 self, param_dtype, run_checks, *fsdp_args, **fsdp_kwargs
268 self, param_dtype, run_checks, *fsdp_args, **fsdp_kwargs
282 def _get_simple_model(self, param_dtype, *fsdp_args, **fsdp_kwargs):
|
| D | test_fsdp_checkpoint.py | 83 **fsdp_kwargs,
|
| /external/pytorch/torch/testing/_internal/ |
| D | common_fsdp.py | 391 **fsdp_kwargs,
519 **fsdp_kwargs,
718 **fsdp_kwargs,
897 **fsdp_kwargs,
1302 **fsdp_kwargs,
|