Skip to content

Commit 71191a6

Browse files
authored
Clean up cuTENSOR plan memory (#212)
1 parent 541f2f4 commit 71191a6

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

ext/TensorOperationscuTENSORExt.jl

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -312,7 +312,10 @@ function plan_trace(@nospecialize(A::AbstractArray), Ainds::ModeType,
312312
plan_pref = Ref{cutensorPlanPreference_t}()
313313
cutensorCreatePlanPreference(handle(), plan_pref, algo, jit)
314314

315-
return CuTensorPlan(desc[], plan_pref[]; workspacePref=workspace)
315+
plan = CuTensorPlan(desc[], plan_pref[]; workspacePref=workspace)
316+
cuTENSOR.cutensorDestroyOperationDescriptor(desc[])
317+
cuTENSOR.cutensorDestroyPlanPreference(plan_pref[])
318+
return plan
316319
end
317320

318321
end

0 commit comments

Comments
 (0)