[AutoParallel] Fix bug of return PyNone. (#60773)

PaddlePaddle · Jan 15, 2024 · 53acf64 · 53acf64
1 parent bc0b16c
commit 53acf64
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -2464,7 +2464,10 @@ paddle::Tensor PyTensorHook::operator()(const paddle::Tensor& var) {
 
   PyObject* res = nullptr;
   try {
-    bool return_py_none_if_not_initialize = var.is_dist_tensor() ? false : true;
+    bool return_py_none_if_not_initialize = true;
+    if (var.defined() && !var.initialized()) {
+      return_py_none_if_not_initialize = !var.is_dist_tensor();
+    }
     PyObject* p_tmp_var = ToPyObject(var, return_py_none_if_not_initialize);
     res = PyObject_CallFunctionObjArgs(py_func_, p_tmp_var, nullptr);
     Py_DECREF(p_tmp_var);

diff --git a/python/paddle/distributed/fleet/utils/mix_precision_utils.py b/python/paddle/distributed/fleet/utils/mix_precision_utils.py
@@ -58,7 +58,7 @@ def param_hook(tmp_grad):
             ), "In main_grad node, param.grad should be None, but find param[{}] has grad.".format(
                 param.name
             )
-            if tmp_grad._is_initialized():
+            if tmp_grad is not None and tmp_grad._is_initialized():
                 # Some previous pylayer may return None, should check grad validation.
                 if param.main_grad is None:
                     param.main_grad = core.eager.Tensor(