From 3c0bd3afbb4349424688f3adc62681ab4cca3bd2 Mon Sep 17 00:00:00 2001 From: wanghuancoder Date: Tue, 22 Nov 2022 10:52:17 +0800 Subject: [PATCH] use full directly if device is CPU and in dygraph, for optimizer (#48189) --- python/paddle/optimizer/optimizer.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/python/paddle/optimizer/optimizer.py b/python/paddle/optimizer/optimizer.py index 59663bb8190..6a9d504cc29 100644 --- a/python/paddle/optimizer/optimizer.py +++ b/python/paddle/optimizer/optimizer.py @@ -730,10 +730,22 @@ class Optimizer: ) if device is None: device = self._get_device_for_param(param.name) - with device_guard(device): - self.helper.set_variable_initializer( - var, initializer=Constant(value=float(fill_value)) + + if in_dygraph_mode() and ( + device == 'cpu' or isinstance(device, core.CPUPlace) + ): + _C_ops.full_( + var, + var.shape, + str(float(fill_value)), + var.dtype, + core.CPUPlace(), ) + else: + with device_guard(device): + self.helper.set_variable_initializer( + var, initializer=Constant(value=float(fill_value)) + ) if framework._non_static_mode(): if len(self._accumulators_holder) > 0: -- GitLab