refactor SGDUpdate and MomentumUpdate UserOp (#4930)
* refactor(ModelUpdate): SGDUpdate and MomentumUpdate use optional input
for learning rate
* fix(*): fix bugs
* style(*): refine code
Co-authored-by:
oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>
Showing
- oneflow/python/nn/optimizer/sgd.py 3 additions, 7 deletionsoneflow/python/nn/optimizer/sgd.py
- oneflow/user/kernels/model_update_kernel_util.cpp 16 additions, 16 deletionsoneflow/user/kernels/model_update_kernel_util.cpp
- oneflow/user/kernels/model_update_kernel_util.cu 39 additions, 35 deletionsoneflow/user/kernels/model_update_kernel_util.cu
- oneflow/user/kernels/model_update_kernel_util.h 5 additions, 4 deletionsoneflow/user/kernels/model_update_kernel_util.h
- oneflow/user/kernels/model_update_kernels.cpp 17 additions, 4 deletionsoneflow/user/kernels/model_update_kernels.cpp
- oneflow/user/ops/model_update_ops.cpp 12 additions, 6 deletionsoneflow/user/ops/model_update_ops.cpp
Please register or sign in to comment