Unverified Commit f3d77f6c authored by jberchtold-nvidia's avatar jberchtold-nvidia Committed by GitHub
Browse files

[JAX] Fix NVTETensor leak in attention.cpp (#1841)



Fix NVTETensor leak in attention.cpp
Signed-off-by: default avatarJeremy Berchtold <jberchtold@nvidia.com>
parent c1417116
...@@ -199,6 +199,8 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( ...@@ -199,6 +199,8 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes(
} }
} }
nvte_tensor_pack_destroy(&aux_output_tensors);
auto workspace_shape = MakeShapeVector(query_workspace_tensor.shape()); auto workspace_shape = MakeShapeVector(query_workspace_tensor.shape());
return pybind11::make_tuple(workspace_shape, query_workspace_tensor.dtype()); return pybind11::make_tuple(workspace_shape, query_workspace_tensor.dtype());
} }
...@@ -485,6 +487,8 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( ...@@ -485,6 +487,8 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes(
} }
} }
nvte_tensor_pack_destroy(&aux_input_tensors);
auto work_shape = MakeShapeVector(query_workspace_tensor.shape()); auto work_shape = MakeShapeVector(query_workspace_tensor.shape());
return pybind11::make_tuple(work_shape, query_workspace_tensor.dtype()); return pybind11::make_tuple(work_shape, query_workspace_tensor.dtype());
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment