Searched refs:saved_tensors_hooks (Results 1 – 9 of 9) sorted by relevance
/aosp_15_r20/external/pytorch/torch/distributed/tensor/parallel/ |
H A D | input_reshard.py | 42 cx: Optional[torch.autograd.graph.saved_tensors_hooks] = None 45 saved_tensor_hooks = torch.autograd.graph.saved_tensors_hooks(
|
/aosp_15_r20/external/pytorch/torch/autograd/ |
H A D | graph.py | 208 class saved_tensors_hooks: class 287 class save_on_cpu(saved_tensors_hooks): 575 class _swap_with_cloned(saved_tensors_hooks):
|
/aosp_15_r20/external/pytorch/test/distributed/fsdp/ |
H A D | test_checkpoint_wrapper.py | 356 orig_init = torch.autograd.graph.saved_tensors_hooks.__init__ 357 torch.autograd.graph.saved_tensors_hooks.__init__ = patched_init 386 torch.autograd.graph.saved_tensors_hooks.__init__ = orig_init
|
/aosp_15_r20/external/pytorch/test/ |
H A D | test_autograd.py | 568 ctx_1 = torch.autograd.graph.saved_tensors_hooks(lambda x: x, unpack) 569 ctx_2 = torch.autograd.graph.saved_tensors_hooks(lambda x: x, lambda x: x) 593 ctx_1 = torch.autograd.graph.saved_tensors_hooks(lambda x: x, unpack) 1946 class RegisterPostNodeHook(torch.autograd.graph.saved_tensors_hooks): 9410 with torch.autograd.graph.saved_tensors_hooks(pack, lambda x: x): 9419 with torch.autograd.graph.saved_tensors_hooks(lambda x: x, lambda x: x): 9427 with torch.autograd.graph.saved_tensors_hooks(lambda x: 2 * x, lambda x: x / 2): 9435 with torch.autograd.graph.saved_tensors_hooks(lambda x: 2 * x, lambda x: x): 9456 with torch.autograd.graph.saved_tensors_hooks(lambda x: x, lambda x: x): 9464 with torch.autograd.graph.saved_tensors_hooks(lambda x: x, lambda x: x): [all …]
|
/aosp_15_r20/external/pytorch/torch/utils/ |
H A D | checkpoint.py | 1051 class _recomputation_hook(torch.autograd.graph.saved_tensors_hooks): 1099 class _checkpoint_hook(torch.autograd.graph.saved_tensors_hooks):
|
/aosp_15_r20/external/pytorch/docs/source/notes/ |
H A D | autograd.rst | 783 :class:`~torch.autograd.graph.saved_tensors_hooks` to register a pair of 808 with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook): 827 with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook): 836 with torch.autograd.graph.saved_tensors_hooks(lambda x: x, lambda x: x):
|
/aosp_15_r20/external/pytorch/docs/source/ |
H A D | autograd.rst | 353 .. autoclass:: torch.autograd.graph.saved_tensors_hooks
|
/aosp_15_r20/external/pytorch/test/inductor/ |
H A D | test_compiled_autograd.py | 2660 with torch.autograd.graph.saved_tensors_hooks(
|
/aosp_15_r20/external/pytorch/test/dynamo/ |
H A D | test_repros.py | 1104 with torch.autograd.graph.saved_tensors_hooks(pack_hook, unpack_hook):
|