Skip to content

Commit

Permalink
Update op_def_struct.h to fix memory leak (#888)
Browse files Browse the repository at this point in the history
This commit removes the .release() calls, ensuring that CreateLiteCustomOpV2 now returns std::shared_ptr, maintaining proper smart pointer ownership and preventing memory corruption.

Co-authored-by: Wenbing Li <[email protected]>
  • Loading branch information
yunmengxie and wenbingl authored Feb 9, 2025
1 parent 48176ac commit 30609b7
Showing 1 changed file with 4 additions and 4 deletions.
8 changes: 4 additions & 4 deletions include/op_def_struct.h
Original file line number Diff line number Diff line change
Expand Up @@ -286,18 +286,18 @@ struct OrtLiteCustomStructV2 : public OrtLiteCustomOp {
};

template <typename RType, typename... Args>
OrtLiteCustomOp* CreateLiteCustomOpV2(const char* op_name,
std::shared_ptr<OrtLiteCustomOp> CreateLiteCustomOpV2(const char* op_name,
const char* execution_provider,
RType (*custom_compute_fn)(Args...)) {
using LiteOp = OrtLiteCustomStructV2<FunctionKernel<RType, Args...>>;
return std::make_unique<LiteOp>(op_name, execution_provider, custom_compute_fn).release();
return std::make_shared<LiteOp>(op_name, execution_provider, custom_compute_fn);
}

template <typename OpKernel>
OrtLiteCustomOp* CreateLiteCustomOpV2(const char* op_name,
std::shared_ptr<OrtLiteCustomOp> CreateLiteCustomOpV2(const char* op_name,
const char* execution_provider) {
using LiteOp = OrtLiteCustomStructV2<OpKernel>;
return std::make_unique<LiteOp>(op_name, execution_provider).release();
return std::make_shared<LiteOp>(op_name, execution_provider);
}

} // namespace Custom
Expand Down

0 comments on commit 30609b7

Please sign in to comment.