From b3c0543114eff9a621470798d9a4e22cfe04bea3 Mon Sep 17 00:00:00 2001
From: lucasdelimanogueira <lucasln96@hotmail.com>
Date: Thu, 16 May 2024 19:53:39 -0300
Subject: [PATCH] batched matmul cuda

---
 norch/csrc/cuda.cu | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/norch/csrc/cuda.cu b/norch/csrc/cuda.cu
index e79955a..9ebe6d3 100644
--- a/norch/csrc/cuda.cu
+++ b/norch/csrc/cuda.cu
@@ -402,7 +402,7 @@ __host__ void batched_matmul_tensor_cuda(Tensor* tensor1, Tensor* tensor2, float
     int cols2 = tensor2->shape[2];
 
     dim3 threadsPerBlock(16, 16);
-    dim3 number_of_blocks((cols2 + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows1 + threadsPerBlock.y - 1) / threadsPerBlock.y);
+    dim3 number_of_blocks((cols2 + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows1 + threadsPerBlock.y - 1) / threadsPerBlock.y, batch_size);
     batched_matmul_tensor_cuda_kernel<<<number_of_blocks, threadsPerBlock>>>(tensor1->data, tensor2->data, result_data, batch_size, rows1, cols1, cols2);