diff --git a/src/layers/conv2d.cu b/src/layers/conv2d.cu
index bbcc516..b7e83a3 100644
--- a/src/layers/conv2d.cu
+++ b/src/layers/conv2d.cu
@@ -26,11 +26,12 @@ Conv2d::Conv2d(
       stride(stride),
       numFilters(numFilters),
       paddingSize(paddingSize) {
+
     outputSize = {
         (inputSize.first - kernelSize.first + 2 * paddingSize.first) /
                 stride.first + 1,
-        (inputSize.first - kernelSize.first + 2 * paddingSize.first) /
-                stride.first + 1
+        (inputSize.second - kernelSize.second + 2 * paddingSize.second) /
+                stride.second + 1
     };
 
     activation =