未验证 提交 5eb381a3 编写于 作者: Z Zeng Jinle 提交者: GitHub

refine reallocate of workspace size, test=develop (#19843)

上级 71b2ed61
...@@ -218,6 +218,8 @@ class CudnnWorkspaceHandle { ...@@ -218,6 +218,8 @@ class CudnnWorkspaceHandle {
if (required_workspace_bytes <= WorkspaceSize()) { if (required_workspace_bytes <= WorkspaceSize()) {
return; return;
} }
// reset allocation first before re-allocate to save memory
allocation_.reset();
allocation_ = memory::Alloc(device_context_, required_workspace_bytes); allocation_ = memory::Alloc(device_context_, required_workspace_bytes);
} }
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册