From 3ea7fb0f9f8e4d40a0b0dd7f998eb16fbb379d7d Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Fri, 14 Jan 2022 14:54:19 +0800 Subject: [PATCH] fix mem leak (#10272) --- .../cpu/aten_op_executor/aten_op_executor.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/aten_op_executor.cc b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/aten_op_executor.cc index 9ffbb7961b..6e44430bb8 100644 --- a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/aten_op_executor.cc +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/aten_op_executor.cc @@ -68,9 +68,11 @@ struct ATenOperator { bool is_list = is_list_arguments[index]; c10::IValue i_value; + // Create the torch tensor from this DLPack no matter we need it or not below, + // so that the dlpack's deleter will be triggered when torch tensor is out of scope. + at::Tensor tensor = at::fromDLPack(dlpack); switch (elem_kinds[index]) { case c10::TypeKind::TensorType: { - at::Tensor tensor = at::fromDLPack(dlpack); i_value = is_optional ? c10::IValue(c10::optional(tensor)) : c10::IValue(tensor); } break; case c10::TypeKind::IntType: {