1 files changed, 5 insertions, 2 deletions
diff --git a/src/core/NEON/kernels/assembly/NEGEMMInterleavedPrepareBWrapperKernel.cpp b/src/core/NEON/kernels/assembly/NEGEMMInterleavedPrepareBWrapperKernel.cpp
index 6c201cedb3..41a031c1c7 100644
--- a/src/core/NEON/kernels/assembly/NEGEMMInterleavedPrepareBWrapperKernel.cpp
+++ b/src/core/NEON/kernels/assembly/NEGEMMInterleavedPrepareBWrapperKernel.cpp
@@ -68,7 +68,7 @@ void for_each_element_in_window(const Window &window, const ITensor *b, ITensor
 
 // Calculate the size of transformed_b:
 template <typename To, bool use_dot>
-unsigned int get_B_pretransposed_array_size(unsigned int N, unsigned int K, const BlockSizes &bs)
+unsigned int get_B_pretransposed_array_size(unsigned int N, unsigned int K, const BlockSizes &bs, unsigned int multis)
 {
     using strategy = typename Kernel<To, use_dot>::strategy;
 
@@ -89,6 +89,9 @@ unsigned int get_B_pretransposed_array_size(unsigned int N, unsigned int K, cons
     // Calculate the total size of the buffer:
     size_t total = num_full_k * normal_k_size * (num_full_x * normal_x_size + left_over_x_size);
     total += left_over_k_size * (left_over_x_size + num_full_x * normal_x_size);
+
+    total *= multis;
+
     return total;
 }
 
@@ -114,7 +117,7 @@ void NEGEMMInterleavedPrepareBWrapperKernelTemplate<To, use_dot>::configure(cons
 
     _block_sizes = calculate_block_sizes<strategy>(ci, params.M, params.N, params.K);
 
-    auto_init_if_empty(*transformed_b->info(), b->info()->clone()->set_tensor_shape(TensorShape{ get_B_pretransposed_array_size<To, use_dot>(_Nsize, _Ksize, _block_sizes) }));
+    auto_init_if_empty(*transformed_b->info(), b->info()->clone()->set_tensor_shape(TensorShape{ get_B_pretransposed_array_size<To, use_dot>(_Nsize, _Ksize, _block_sizes, multis) }));
 
     Window window;
     window.set(Window::DimX, Window::Dimension(0, ceil_to_multiple(_Nsize, _block_sizes.x_block), _block_sizes.x_block));