add batch size to the mkldnn matmul cache key (#24408)

test=develop
release/2.0-alpha
Wojciech Uss 6 years ago committed by GitHub
parent 53125c2f6f
commit fddf442470
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -322,9 +322,10 @@ static std::shared_ptr<MatMulFactory<XT, YT, OT>> GetPrimitiveFactory(
const ExecutionContext& ctx) {
const auto& out_name = ctx.OutputName("Out");
const auto& dev_ctx = ctx.template device_context<MKLDNNDeviceContext>();
const auto batch_size = ctx.Input<Tensor>("X")->dims()[0];
const std::string key =
platform::CreateKey(platform::ThreadIDasStr(), out_name);
platform::CreateKey(platform::ThreadIDasStr(), batch_size, out_name);
auto factory =
std::static_pointer_cast<MatMulFactory<XT, YT, OT>>(dev_ctx.GetBlob(key));

Loading…
Cancel
Save