diff --git a/include/layers/max_pool.hpp b/include/layers/max_pool.hpp
index 0adafe9..515f226 100644
--- a/include/layers/max_pool.hpp
+++ b/include/layers/max_pool.hpp
@@ -8,7 +8,7 @@ class MaxPool2d : public Layer {
   public:
     MaxPool2d(
         CUDANet::Shape        input_shape,
-        CUDANet::Shape        pooling_shape,
+        CUDANet::Shape        pool_shape,
         CUDANet::Shape        stride_shape,
         CUDANet::Shape        padding_shape,
         CUDANet::Backend* backend
@@ -38,7 +38,7 @@ class MaxPool2d : public Layer {
   private:
     CUDANet::Shape in_shape;
 
-    CUDANet::Shape pooling_shape;
+    CUDANet::Shape pool_shape;
     CUDANet::Shape stride_shape;
     CUDANet::Shape padding_shape;
 
diff --git a/include/shape.hpp b/include/shape.hpp
index 901350c..cc39128 100644
--- a/include/shape.hpp
+++ b/include/shape.hpp
@@ -6,4 +6,21 @@ namespace CUDANet {
 
 typedef std::vector<size_t> Shape;
 
-} // namespace CUDANet
+class InvalidShapeException : public std::runtime_error {
+  public:
+    InvalidShapeException(
+        const std::string& param_name,
+        size_t             expected,
+        size_t             actual
+    )
+        : std::runtime_error(
+              std::format(
+                  "Invalid {} shape. Expected {}, actual {}",
+                  param_name,
+                  expected,
+                  actual
+              )
+          ) {}
+};
+
+}  // namespace CUDANet
diff --git a/src/layers/avg_pooling.cpp b/src/layers/avg_pooling.cpp
index 8bbfe21..7ebd14b 100644
--- a/src/layers/avg_pooling.cpp
+++ b/src/layers/avg_pooling.cpp
@@ -18,35 +18,19 @@ AvgPool2d::AvgPool2d(
       padding_shape(padding_shape),
       backend(backend) {
     if (in_shape.size() != 3) {
-        throw std::runtime_error(
-            std::format(
-                "Invalid input shape. Expected 3 dims, got {}", input_shape.size()
-            )
-        );
+        throw InvalidShapeException("input", 3, in_shape.size());
     }
 
     if (pool_shape.size() != 2) {
-        throw std::runtime_error(
-            std::format(
-                "Invalid pool shape. Expected 2 dims, got {}", pool_shape.size()
-            )
-        );
+        throw InvalidShapeException("pool", 2, pool_shape.size());
     }
 
     if (stride_shape.size() != 2) {
-        throw std::runtime_error(
-            std::format(
-                "Invalid stride shape. Expected 2 dims, got {}", stride_shape.size()
-            )
-        );
+        throw InvalidShapeException("stride", 2, stride_shape.size());
     }
 
     if (padding_shape.size() != 2) {
-        throw std::runtime_error(
-            std::format(
-                "Invalid padding shape. Expected 2 dims, got {}", padding_shape.size()
-            )
-        );
+        throw InvalidShapeException("padding", 2, padding_shape.size());
     }
 
     out_shape = {
diff --git a/src/layers/conv2d.cpp b/src/layers/conv2d.cpp
index b742267..88e71db 100644
--- a/src/layers/conv2d.cpp
+++ b/src/layers/conv2d.cpp
@@ -21,47 +21,31 @@ Conv2d::Conv2d(
       padding_shape(padding_shape),
       backend(backend) {
     if (in_shape.size() != 3) {
-        throw std::runtime_error(
-            std::format(
-                "Invalid input shape. Expected 3 dims, got {}", in_shape.size()
-            )
-        );
+        throw InvalidShapeException("input", 3, in_shape.size());
     }
 
     if (kernel_shape.size() != 3) {
-        throw std::runtime_error(
-            std::format(
-                "Invalid kernel shape. Expected 3 dims, got {}", kernel_shape.size()
-            )
-        );
+        throw InvalidShapeException("kernel", 3, kernel_shape.size());
     }
 
     if (stride_shape.size() != 2) {
-        throw std::runtime_error(
-            std::format(
-                "Invalid stride shape. Expected 2 dims, got {}", stride_shape.size()
-            )
-        );
+        throw InvalidShapeException("stride", 3, stride_shape.size());
     }
 
     if (padding_shape.size() != 2) {
-        throw std::runtime_error(
-            std::format(
-                "Invalid padding shape. Expected 2 dims, got {}", padding_shape.size()
-            )
-        );
+        throw InvalidShapeException("padding", 3, padding_shape.size());
     }
 
-    size_t out_h = (in_shape[0] - kernel_shape[0] + 2 * padding_shape[0]) /
-                       stride_shape[0] +
-                   1;
-    size_t out_w = (in_shape[1] - kernel_shape[1] + 2 * padding_shape[1]) /
-                       stride_shape[1] +
-                   1;
-    out_shape.resize(3);
-    out_shape[0] = out_h;
-    out_shape[1] = out_w;
-    out_shape[2] = kernel_shape[2];
+    out_shape = {
+        (in_shape[0] - kernel_shape[0] + 2 * padding_shape[0]) /
+                stride_shape[0] +
+            1,
+        (in_shape[1] - kernel_shape[1] + 2 * padding_shape[1]) /
+                stride_shape[1] +
+            1,
+        kernel_shape[2]
+    };
+
     output = CUDANet::Tensor(
         Shape{out_shape[0] * out_shape[1] * out_shape[3]},
         CUDANet::DType::FLOAT32, backend
@@ -69,7 +53,7 @@ Conv2d::Conv2d(
 
     weights = CUDANet::Tensor(
         Shape{
-            kernel_shape[0] * kernel_shape[1] * kernel_shape[2] * in_shape[2]
+            kernel_shape[0], kernel_shape[1], kernel_shape[2], in_shape[2]
         },
         CUDANet::DType::FLOAT32, backend
     );
@@ -83,18 +67,11 @@ Conv2d::Conv2d(
 
 Conv2d::~Conv2d() {}
 
-CUDANet::Tensor& Conv2d::forward( CUDANet::Tensor& input) {
+CUDANet::Tensor& Conv2d::forward(CUDANet::Tensor& input) {
     output.zero();
     backend->conv2d(
-        weights,
-        biases,
-        input,
-        output,
-        in_shape,
-        padding_shape,
-        kernel_shape,
-        stride_shape,
-        out_shape
+        weights, biases, input, output, in_shape, padding_shape, kernel_shape,
+        stride_shape, out_shape
     );
     return output;
 }
diff --git a/src/layers/dense.cpp b/src/layers/dense.cpp
index 57cf5e2..b7f6c77 100644
--- a/src/layers/dense.cpp
+++ b/src/layers/dense.cpp
@@ -5,26 +5,22 @@
 
 using namespace CUDANet::Layers;
 
-Dense::Dense(CUDANet::Shape in, CUDANet::Shape out, CUDANet::Backend* backend)
+Dense::Dense(CUDANet::Shape in_shape, CUDANet::Shape out_shape, CUDANet::Backend* backend)
     : backend(backend),
-      in_shape(in),
-      out_shape(out) {
+      in_shape(in_shape),
+      out_shape(out_shape) {
 
-    if (in.size() != 1) {
-        throw std::runtime_error(
-            std::format("Invalid shape. Expected [1], got {}", in_shape)
-        );
+    if (in_shape.size() != 1) {
+        throw InvalidShapeException("input", 1, in_shape.size());
     }
 
-    if (out.size() != 1) {
-        throw std::runtime_error(
-            std::format("Invalid shape. Expected [1], got {}", out_shape)
-        );
+    if (out_shape.size() != 1) {
+        throw InvalidShapeException("output", 1, out_shape.size());
     }
 
-    weights = CUDANet::Tensor(Shape{in[0] * out[0]}, CUDANet::DType::FLOAT32, backend);
-    biases = CUDANet::Tensor(Shape{out[0]}, CUDANet::DType::FLOAT32, backend);
-    output = CUDANet::Tensor(Shape{out[0]}, CUDANet::DType::FLOAT32, backend);
+    weights = CUDANet::Tensor(Shape{out_shape[0], in_shape[0]}, CUDANet::DType::FLOAT32, backend);
+    biases = CUDANet::Tensor(Shape{out_shape[0]}, CUDANet::DType::FLOAT32, backend);
+    output = CUDANet::Tensor(Shape{out_shape[0]}, CUDANet::DType::FLOAT32, backend);
 
     weights.zero();
     biases.zero();
diff --git a/src/layers/max_pool.cpp b/src/layers/max_pool.cpp
index 20a7666..b56ffac 100644
--- a/src/layers/max_pool.cpp
+++ b/src/layers/max_pool.cpp
@@ -6,27 +6,39 @@ using namespace CUDANet::Layers;
 
 MaxPool2d::MaxPool2d(
     CUDANet::Shape    input_shape,
-    CUDANet::Shape    pooling_shape,
+    CUDANet::Shape    pool_shape,
     CUDANet::Shape    stride_shape,
     CUDANet::Shape    padding_shape,
     CUDANet::Backend* backend
 )
     : in_shape(input_shape),
-      pooling_shape(pooling_shape),
+      pool_shape(pool_shape),
       stride_shape(stride_shape),
       padding_shape(padding_shape),
       backend(backend) {
-    size_t out_h = (in_shape[0] + 2 * padding_shape[0] - pooling_shape[0]) /
-                       stride_shape[0] +
-                   1;
-    size_t out_w = (in_shape[1] + 2 * padding_shape[1] - pooling_shape[1]) /
-                       stride_shape[1] +
-                   1;
+    if (in_shape.size() != 3) {
+        throw InvalidShapeException("input", 3, in_shape.size());
+    }
 
-    out_shape.resize(3);
-    out_shape[0] = out_h;
-    out_shape[1] = out_w;
-    out_shape[2] = in_shape[2];
+    if (pool_shape.size() != 2) {
+        throw InvalidShapeException("pool", 2, pool_shape.size());
+    }
+
+    if (stride_shape.size() != 2) {
+        throw InvalidShapeException("stride", 2, stride_shape.size());
+    }
+
+    if (padding_shape.size() != 2) {
+        throw InvalidShapeException("padding", 2, padding_shape.size());
+    }
+
+    out_shape = {
+        (in_shape[0] + 2 * padding_shape[0] - pool_shape[0]) / stride_shape[0] +
+            1,
+        (in_shape[1] + 2 * padding_shape[1] - pool_shape[1]) / stride_shape[1] +
+            1,
+        in_shape[2]
+    };
 
     output = CUDANet::Tensor(
         Shape{out_shape[0] * out_shape[1] * out_shape[3]},
@@ -39,7 +51,7 @@ MaxPool2d::~MaxPool2d() {}
 CUDANet::Tensor& MaxPool2d::forward(CUDANet::Tensor& input) {
     output.zero();
     backend->maxPool2d(
-        input, output, in_shape, pooling_shape, stride_shape, padding_shape,
+        input, output, in_shape, pool_shape, stride_shape, padding_shape,
         out_shape
     );
     return output;