Skip to content

GitLab

Unverified Commit 350d64f3 authored by Qianfeng's avatar Qianfeng Committed by GitHub
Browse files

Add workspace setting up for batchnorm bwd/fwd client examples (#860)

parent 7c71dc7e
Show whitespace changes
Inline Side-by-side
Showing with 12 additions and 0 deletions
client_example/13_batchnorm/batchnorm_bwd_nhwc.cpp
View file @ 350d64f3
...@@ -191,6 +191,12 @@ int main(int argc, char* argv[]) ...@@ -191,6 +191,12 @@ int main(int argc, char* argv[])
if(op_ptr->IsSupportedArgument(argument_ptr.get())) if(op_ptr->IsSupportedArgument(argument_ptr.get()))
{ {
size_t workspace_sz = op_ptr->GetWorkSpaceSize(argument_ptr.get());
SimpleDeviceMem workspace(workspace_sz);
op_ptr->SetWorkSpacePointer(argument_ptr.get(), workspace.GetDeviceBuffer());
invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false}); invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false});
} }
......
client_example/13_batchnorm/batchnorm_fwd_nhwc.cpp
View file @ 350d64f3
...@@ -187,6 +187,12 @@ int main(int argc, char* argv[]) ...@@ -187,6 +187,12 @@ int main(int argc, char* argv[])
if(op_ptr->IsSupportedArgument(argument_ptr.get())) if(op_ptr->IsSupportedArgument(argument_ptr.get()))
{ {
size_t workspace_sz = op_ptr->GetWorkSpaceSize(argument_ptr.get());
SimpleDeviceMem workspace(workspace_sz);
op_ptr->SetWorkSpacePointer(argument_ptr.get(), workspace.GetDeviceBuffer());
invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false}); invoker_ptr->Run(argument_ptr.get(), StreamConfig{nullptr, false});
} }
......
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