Home
last modified time | relevance | path

Searched refs:stack_module_state (Results 1 – 6 of 6) sorted by relevance

/aosp_15_r20/external/pytorch/functorch/examples/ensembling/
H A Dparallel_train.py7 from torch.func import functional_call, grad_and_value, stack_module_state, vmap
118 params, _ = stack_module_state(models)
/aosp_15_r20/external/pytorch/docs/source/
H A Dfunc.migrating.rst157 Please use :func:`torch.func.stack_module_state` instead of
159 :func:`torch.func.stack_module_state` returns two dictionaries, one of stacked parameters, and
190 params, buffers = torch.func.stack_module_state(models)
H A Dfunc.api.rst69 stack_module_state
/aosp_15_r20/external/pytorch/torch/func/
H A D__init__.py11 from torch._functorch.functional_call import functional_call, stack_module_state
/aosp_15_r20/external/pytorch/torch/_functorch/
H A Dfunctional_call.py159 def stack_module_state( function
/aosp_15_r20/external/pytorch/test/functorch/
H A Dtest_eager_transforms.py52 from torch.func import functional_call, linearize, stack_module_state
3949 _ = stack_module_state(models)
3956 params, buffers = stack_module_state(models)
3968 params, buffers = stack_module_state(models)
3978 _ = stack_module_state(models)
3986 _ = stack_module_state(models)
3993 _ = stack_module_state(models)
4426 return stack_module_state(models)[0]
4527 return stack_module_state(models)[0]