Fix small layer issues

src/backends/cuda/layer_ops.cu

@@ -98,7 +98,7 @@ CUDANet::Tensor& CUDA::conv2d(
     dim3 grid(
         (out_shape[0] + block.x - 1) / block.x,
         (out_shape[1] + block.y - 1) / block.y,
-        (out_shape[3] + block.z - 1) / block.z
+        (out_shape[2] + block.z - 1) / block.z
     );
 
     Kernels::convolution<<<grid, block>>>(
@@ -212,6 +212,7 @@ CUDANet::Tensor& CUDA::batch_norm(
         CUDA_CHECK(cudaGetLastError());
     }
     CUDA_CHECK(cudaDeviceSynchronize());
+    return output;
 }
 
 CUDANet::Tensor& CUDA::concat(

src/layers/add.cpp

@@ -19,4 +19,10 @@ Add::~Add() {}
 CUDANet::Tensor&
 Add::forward(CUDANet::Tensor& input_a, CUDANet::Tensor& input_b) {
     output.zero();
+    backend->add(
+        input_a,
+        input_b,
+        output
+    );
+    return output;
 }

src/layers/avg_pooling.cpp

@@ -76,7 +76,7 @@ size_t AvgPool2d::input_size() {
 }
 
 size_t AvgPool2d::output_size() {
-    return sizeof(float) * out_shape[0] * out_shape[1] * out_shape[3];
+    return sizeof(float) * out_shape[0] * out_shape[1] * out_shape[2];
 }
 
 void AvgPool2d::set_weights(void* input) {}

src/layers/conv2d.cpp

@@ -47,7 +47,7 @@ Conv2d::Conv2d(
     };
 
     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
     );
 

src/layers/dense.cpp

@@ -30,6 +30,7 @@ Dense::Dense(CUDANet::Shape in_shape, CUDANet::Shape out_shape, CUDANet::Backend
 Dense::~Dense() {}
 
 CUDANet::Tensor& Dense::forward(CUDANet::Tensor& input) {
+    output.zero();
     backend->dense(weights, biases, input, output, in_shape[0], out_shape[0]);
     return output;
 }

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
     );
 }