[DoubleGrad PR #6] Fixed issues with TensorWrapper::recover() interfa…

…ce (PaddlePaddle#41287)
AnnaTrainingG · Apr 2, 2022 · a5e00bb · a5e00bb
1 parent f482613
commit a5e00bb
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 8 deletions.
diff --git a/paddle/fluid/eager/auto_code_generator/final_state_generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/final_state_generator/eager_gen.py
@@ -1249,9 +1249,9 @@ def GenerateNodeDefinition(self, grad_node_creation_str):
 
             is_optional = (name in self.optional_inputs)
             if is_optional:
-                tensor_wrapper_recover_str = f"{indent}auto {transformed_tensor_name} = egr::EagerUtils::RecoverOptionalTensorWrapper(&this->{tensor_wrapper_name}, nullptr);"
+                tensor_wrapper_recover_str = f"{indent}auto {transformed_tensor_name} = egr::EagerUtils::RecoverOptionalTensorWrapper(&this->{tensor_wrapper_name}, this->shared_from_this());"
             else:
-                tensor_wrapper_recover_str = f"{indent}auto {transformed_tensor_name} = egr::EagerUtils::RecoverTensorWrapper(&this->{tensor_wrapper_name}, nullptr);"
+                tensor_wrapper_recover_str = f"{indent}auto {transformed_tensor_name} = egr::EagerUtils::RecoverTensorWrapper(&this->{tensor_wrapper_name}, this->shared_from_this());"
             grad_api_args[grad_api_position] = transformed_tensor_name
             get_grad_in_args_list.append(tensor_wrapper_recover_str)
 

diff --git a/paddle/fluid/eager/grad_node_info.h b/paddle/fluid/eager/grad_node_info.h
@@ -87,7 +87,7 @@ class GradSlotMeta {
   std::shared_ptr<phi::DenseTensorMeta> meta_ = nullptr;
 };
 
-class GradNodeBase {
+class GradNodeBase : public std::enable_shared_from_this<GradNodeBase> {
  public:
   GradNodeBase() { VLOG(6) << "Construct GradNodeBase"; }
   GradNodeBase(size_t bwd_in_slot_num, size_t bwd_out_slot_num);

diff --git a/paddle/fluid/eager/tensor_wrapper.h b/paddle/fluid/eager/tensor_wrapper.h
@@ -95,18 +95,19 @@ class TensorWrapper {
     }
 
     check_inplace_version();
+
     // if it's full_reserved just return the full copy of tensor
-    if (full_reserved_) {
-      return intermidiate_tensor_;
-    } else {
+    paddle::experimental::Tensor recovered_tensor = intermidiate_tensor_;
+    if (!full_reserved_) {
       std::shared_ptr<GradNodeBase> new_grad_node = grad_node;
       auto p_ab_autograd_meta =
           std::make_shared<AutogradMeta>(Edge(new_grad_node, out_rank_info_));
-      intermidiate_tensor_.set_autograd_meta(
+      recovered_tensor.set_autograd_meta(
           std::static_pointer_cast<paddle::experimental::AbstractAutogradMeta>(
               p_ab_autograd_meta));
-      return intermidiate_tensor_;
     }
+
+    return recovered_tensor;
   }
 
   void check_inplace_version() {