p04896132
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!778 fix bug of ApplyRMSProp 

Merge pull request !778 from fary86/fix_bug_of_ApplyRMSProp
This commit is contained in:
mindspore-ci-bot 2020-04-28 09:51:41 +08:00 committed by Gitee
parent 0fbec0f55a 78b4479ac6
commit 1b5fb395cc
1 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
mindspore/ops/operations/nn_ops.py
View File

@ -1586,9 +1586,11 @@ class ApplyRMSProp(PrimitiveWithInfer):
args = {"var": var_dtype, "mean_square": mean_square_dtype, "moment": moment_dtype, "grad": grad_dtype}
validator.check_tensor_type_same(args, mstype.number_type, self.name)
args = {"learning_rate": learning_rate_dtype, "decay": decay_dtype,
'momentum': momentum_dtype, "epsilon": epsilon_dtype}
validator.check_scalar_or_tensor_type_same(args, [mstype.float16, mstype.float32], self.name)
valid_types = [mstype.float16, mstype.float32]
args_decay = {"decay": decay_dtype, 'momentum': momentum_dtype, "epsilon": epsilon_dtype}
validator.check_type_same(args_decay, valid_types, self.name)
args_lr = {"learning_rate": learning_rate_dtype, "decay": decay_dtype}
validator.check_scalar_or_tensor_type_same(args_lr, valid_types, self.name, allow_mix=True)
return var_dtype