Browse Source

update definition of operator ApplyMomentumD

tags/v0.3.0-alpha
yanghaoran 5 years ago
parent
commit
066dff61fc
1 changed files with 12 additions and 0 deletions
  1. +12
    -0
      third_party/fwkacllib/inc/ops/nn_training_ops.h

+ 12
- 0
third_party/fwkacllib/inc/ops/nn_training_ops.h View File

@@ -299,6 +299,18 @@ REG_OP(ApplyMomentumCCE)
.ATTR(use_locking, Bool, false) .ATTR(use_locking, Bool, false)
.OP_END_FACTORY_REG(ApplyMomentumCCE) .OP_END_FACTORY_REG(ApplyMomentumCCE)


REG_OP(ApplyMomentumD)
.INPUT(var, TensorType::NumberType())
.INPUT(accum, TensorType::NumberType())
.INPUT(lr, TensorType::NumberType())
.INPUT(grad, TensorType::NumberType())
.INPUT(momentum, TensorType::NumberType())
.OUTPUT(var, TensorType::NumberType())
.OUTPUT(accum, TensorType::NumberType())
.ATTR(use_nesterov, Bool, false)
.ATTR(use_locking, Bool, false)
.OP_END_FACTORY_REG(ApplyMomentumD)

/** /**
*@brief Updates "var" according to the AddSign update.\n *@brief Updates "var" according to the AddSign update.\n
* t-1 mean previous period. * t-1 mean previous period.


Loading…
Cancel
Save