Fix small layer issues

2025-12-22 22:34:22 +00:00 · 2025-11-22 00:33:51 +01:00
parent 4c8b2ef537
commit 104d6ea33d
6 changed files with 12 additions and 4 deletions
--- a/src/layers/max_pool.cpp
+++ b/src/layers/max_pool.cpp
@@ -41,7 +41,7 @@ MaxPool2d::MaxPool2d(
    };

    output = CUDANet::Tensor(
-        Shape{out_shape[0] * out_shape[1] * out_shape[3]},
+        Shape{out_shape[0] * out_shape[1] * out_shape[2]},
        CUDANet::DType::FLOAT32, backend
    );
 }