Commit 86449091 authored by wangshaojie6's avatar wangshaojie6
Browse files

use datatype and split k to check whether a workspace is used

parent 53f06c50
...@@ -299,7 +299,7 @@ int main(int argc, char* argv[]) ...@@ -299,7 +299,7 @@ int main(int argc, char* argv[])
// alloc work space // alloc work space
size_t bwd_weight_workspace_size = conv->GetWorkSpaceSize(argument.get()); size_t bwd_weight_workspace_size = conv->GetWorkSpaceSize(argument.get());
float ave_time = 0.f; float ave_time = 0.f;
if(bwd_weight_workspace_size > 0) if(std::is_same<InDataType, ck::bhalf_t>::value && split_k > 1)
{ {
DeviceMem wei_work_space_device_buf(bwd_weight_workspace_size); DeviceMem wei_work_space_device_buf(bwd_weight_workspace_size);
wei_work_space_device_buf.SetZero(); wei_work_space_device_buf.SetZero();
......
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