Home
last modified time | relevance | path

Searched defs:momentum_buffer (Results 1 – 4 of 4) sorted by relevance

/aosp_15_r20/external/executorch/extension/training/optimizer/
H A Dsgd.h43 explicit SGDParamState(executorch::aten::Tensor& momentum_buffer) in SGDParamState()
46 executorch::aten::Tensor& momentum_buffer() { in momentum_buffer() function
/aosp_15_r20/external/pytorch/aten/src/ATen/native/cpu/
H A DFusedSGDKernel.cpp188 const at::Tensor& momentum_buffer, in sgd_fused_step_impl()
238 const at::Tensor& momentum_buffer, in fused_sgd_kernel()
/aosp_15_r20/external/executorch/extension/training/optimizer/test/
H A Dsgd_test.cpp37 Tensor momentum_buffer = tf.make({2, 2}, {1, 2, 3, 4}); in TEST_F() local
/aosp_15_r20/external/pytorch/aten/src/ATen/native/cuda/
H A DFusedSgdKernel.cu42 const auto momentum_buffer = is_first_step in sgd_math() local