diff --git a/python/paddle/incubate/optimizer/distributed_fused_lamb.py b/python/paddle/incubate/optimizer/distributed_fused_lamb.py index f8e3b55aba621..3f3df92be5ebe 100644 --- a/python/paddle/incubate/optimizer/distributed_fused_lamb.py +++ b/python/paddle/incubate/optimizer/distributed_fused_lamb.py @@ -13,13 +13,13 @@ # limitations under the License. import os +import paddle from paddle.fluid import framework, core, layers, unique_name from paddle.fluid.framework import Variable from paddle.fluid.clip import ClipGradByGlobalNorm from paddle.fluid.initializer import Constant from paddle.fluid.layer_helper import LayerHelper from paddle.fluid.optimizer import Optimizer -import paddle.distributed as dist from paddle.distributed.collective import new_group from paddle.fluid.executor import global_scope from paddle.fluid.framework import name_scope @@ -288,8 +288,8 @@ def _apply_gradients_impl(self, params_grads): step = self._get_or_create_step() - rank = dist.get_rank() - nranks = dist.get_world_size() + rank = paddle.distributed.get_rank() + nranks = paddle.distributed.get_world_size() if self._nproc_per_node is None: nproc_per_node = nranks else: