Skip to content

Commit 5f2eebe

Browse files
custom op support post hook (#72976)
1 parent 8e5d148 commit 5f2eebe

File tree

1 file changed

+8
-0
lines changed

1 file changed

+8
-0
lines changed

paddle/fluid/eager/custom_operator/custom_operator_node.cc

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -373,6 +373,10 @@ RunCustomOpNode::operator()(paddle::small_vector<std::vector<paddle::Tensor>,
373373
grad_node->SetAttrs(attrs);
374374
}
375375

376+
if (HasNodePostHook()) {
377+
outs = ApplyNodePostHooks(outs, hooked_grads);
378+
}
379+
376380
return outs;
377381
}
378382

@@ -459,6 +463,10 @@ RunCustomOpDoubleGradNode::operator()(
459463
outs[i] = ctx.OutputsBetween(output_pair.first, output_pair.second);
460464
}
461465

466+
if (HasNodePostHook()) {
467+
outs = ApplyNodePostHooks(outs, hooked_grads);
468+
}
469+
462470
return outs;
463471
}
464472
} // namespace egr

0 commit comments

Comments
 (0)