From d27d7f145318b8aeb8ba3b3b06015aa2b119e466 Mon Sep 17 00:00:00 2001 From: lifuhua Date: Wed, 15 Mar 2023 19:08:32 +0800 Subject: [PATCH] Fix bug: when opt_offload is enabled, the cast operator is not offloaded to the CPU Signed-off-by: lifuhua --- official/nlp/Pangu_alpha/src/adam.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/official/nlp/Pangu_alpha/src/adam.py b/official/nlp/Pangu_alpha/src/adam.py index 4488080e0..5f06c66bd 100644 --- a/official/nlp/Pangu_alpha/src/adam.py +++ b/official/nlp/Pangu_alpha/src/adam.py @@ -39,13 +39,15 @@ def _update_run_kernel(opt, clip_value, beta1, beta2, eps, lr, weight_decay, Update parameters by AdamWeightDecay op. """ success = True + cast = P.Cast() + cast.add_prim_attr("primitive_target", "CPU") if optim_filter: if decay_flags: next_param = opt(param, m, v, lr, beta1, beta2, eps, weight_decay, - P.Cast()(gradient, mstype.float16), clip_value) + cast(gradient, mstype.float16), clip_value) else: next_param = opt(param, m, v, lr, beta1, beta2, eps, 0.0, - P.Cast()(gradient, mstype.float16), clip_value) + cast(gradient, mstype.float16), clip_value) return F.depend(success, next_param) return success -- Gitee