Publishing 2019 R1 content

[platform/upstream/dldt.git] / inference-engine / thirdparty / mkl-dnn / tests / benchdnn / mkldnn_memory.hpp

diff --git a/inference-engine/thirdparty/mkl-dnn/tests/benchdnn/mkldnn_memory.hpp b/inference-engine/thirdparty/mkl-dnn/tests/benchdnn/mkldnn_memory.hpp
index 8c6a4c1..6a1441c 100644 (file)
--- a/inference-engine/thirdparty/mkl-dnn/tests/benchdnn/mkldnn_memory.hpp
+++ b/inference-engine/thirdparty/mkl-dnn/tests/benchdnn/mkldnn_memory.hpp
@@ -77,10 +77,13 @@ struct dnn_mem_t {
 
     size_t size() const { return mkldnn_memory_primitive_desc_get_size(mpd_); }
 
-    size_t nelems() const {
+    size_t nelems(bool with_padding_dims = false) const {
+        auto dims = with_padding_dims
+            ? md_.layout_desc.blocking.padding_dims
+            : md_.dims;
         size_t n = 1;
         for (int i = 0; i < md_.ndims; ++i)
-            n *= md_.dims[i];
+            n *= dims[i];
         return n;
     }