Add getOutputDims to 2d layers

src/layers/avg_pooling.cu

@@ -61,4 +61,8 @@ int AvgPooling2d::getOutputSize() {
 
 int AvgPooling2d::getInputSize() {
     return inputSize.first * inputSize.second * nChannels;
+}
+
+dim2d AvgPooling2d::getOutputDims() {
+    return outputSize;
 }

src/layers/batch_norm.cu

@@ -128,6 +128,10 @@ int BatchNorm2d::getOutputSize() {
     return inputSize.first * inputSize.second * inputChannels;
 }
 
+dim2d BatchNorm2d::getOutputDims() {
+    return inputSize;
+}
+
 float *BatchNorm2d::forward(const float *d_input) {
     // Compute per-channel batch normalization
     for (int i = 0; i < inputChannels; i++) {

src/layers/conv2d.cu

@@ -137,4 +137,8 @@ int Conv2d::getOutputSize() {
 
 int Conv2d::getInputSize() {
     return inputSize.first * inputSize.second * inputChannels;
+}
+
+dim2d Conv2d::getOutputDims() {
+    return outputSize;
 }

src/layers/max_pooling.cu

@@ -59,4 +59,8 @@ int MaxPooling2d::getOutputSize() {
 
 int MaxPooling2d::getInputSize() {
     return inputSize.first * inputSize.second * nChannels;
+}
+
+dim2d MaxPooling2d::getOutputDims() {
+    return outputSize;
 }