diff --git a/examples/inception_v3/inception_v3.cpp b/examples/inception_v3/inception_v3.cpp
index f74ea9c..4ccfb83 100644
--- a/examples/inception_v3/inception_v3.cpp
+++ b/examples/inception_v3/inception_v3.cpp
@@ -4,17 +4,18 @@
 class BasicConv2d : public CUDANet::Module {
   public:
     BasicConv2d(
-        const shape2d      inputSize,
+        const shape2d      inputShape,
         const int          inputChannels,
         const int          outputChannels,
         const shape2d      kernelSize,
         const shape2d      stride,
         const shape2d      padding,
         const std::string &prefix
-    ) {
+    )
+        : outputChannels(outputChannels) {
         // Create the convolution layer
         conv = new CUDANet::Layers::Conv2d(
-            inputSize, inputChannels, kernelSize, stride, outputChannels,
+            inputShape, inputChannels, kernelSize, stride, outputChannels,
             padding, CUDANet::Layers::ActivationType::NONE
         );
 
@@ -25,6 +26,10 @@ class BasicConv2d : public CUDANet::Module {
             CUDANet::Layers::ActivationType::RELU
         );
 
+        inputSize  = inputShape.first * inputShape.second * inputChannels;
+        outputSize = batchNorm->getOutputDims().first *
+                     batchNorm->getOutputDims().second * outputChannels;
+
         addLayer(prefix + ".conv", conv);
         addLayer(prefix + ".bn", batchNorm);
     }
@@ -48,6 +53,8 @@ class BasicConv2d : public CUDANet::Module {
     }
 
   private:
+    int outputChannels;
+
     CUDANet::Layers::Conv2d      *conv;
     CUDANet::Layers::BatchNorm2d *batchNorm;
 };
@@ -55,24 +62,27 @@ class BasicConv2d : public CUDANet::Module {
 class InceptionA : public CUDANet::Module {
   public:
     InceptionA(
-        const shape2d      inputSize,
+        const shape2d      inputShape,
         const int          inputChannels,
         const int          poolFeatures,
         const std::string &prefix
     )
-        : inputSize(inputSize),
+        : inputShape(inputShape),
           inputChannels(inputChannels),
           poolFeatures(poolFeatures) {
+
+        inputSize  = inputShape.first * inputShape.second * inputChannels;
+
         // Branch 1x1
         branch1x1 = new BasicConv2d(
-            inputSize, inputChannels, 64, {1, 1}, {1, 1}, {0, 0},
+            inputShape, inputChannels, 64, {1, 1}, {1, 1}, {0, 0},
             prefix + ".branch1x1"
         );
         addLayer("", branch1x1);
 
         // Branch 5x5
         branch5x5_1 = new BasicConv2d(
-            inputSize, inputChannels, 48, {1, 1}, {1, 1}, {0, 0},
+            inputShape, inputChannels, 48, {1, 1}, {1, 1}, {0, 0},
             prefix + ".branch5x5_1"
         );
         addLayer("", branch5x5_1);
@@ -84,7 +94,7 @@ class InceptionA : public CUDANet::Module {
 
         // Branch 3x3
         branch3x3dbl_1 = new BasicConv2d(
-            inputSize, inputChannels, 64, {1, 1}, {1, 1}, {0, 0},
+            inputShape, inputChannels, 64, {1, 1}, {1, 1}, {0, 0},
             prefix + ".branch3x3dbl_1"
         );
         addLayer("", branch3x3dbl_1);
@@ -101,13 +111,13 @@ class InceptionA : public CUDANet::Module {
 
         // Branch Pool
         branchPool_1 = new CUDANet::Layers::AvgPooling2d(
-            inputSize, inputChannels, {3, 3}, {1, 1}, {1, 1},
+            inputShape, inputChannels, {3, 3}, {1, 1}, {1, 1},
             CUDANet::Layers::ActivationType::NONE
         );
-        addLayer("", branchPool_1);
+        addLayer(prefix + ".branch_pool", branchPool_1);
         branchPool_2 = new BasicConv2d(
             branchPool_1->getOutputDims(), inputChannels, poolFeatures, {1, 1},
-            {1, 1}, {0, 0}, prefix + ".branchPool"
+            {1, 1}, {0, 0}, prefix + ".branch_pool"
         );
         addLayer("", branchPool_2);
 
@@ -121,6 +131,8 @@ class InceptionA : public CUDANet::Module {
         concat_3 = new CUDANet::Layers::Concat(
             concat_2->getOutputSize(), branchPool_2->getOutputSize()
         );
+
+        outputSize = concat_3->getOutputSize();
     }
 
     ~InceptionA() {
@@ -143,15 +155,15 @@ class InceptionA : public CUDANet::Module {
         float *d_branch5x5_out = branch5x5_1->forward(d_input);
         d_branch5x5_out        = branch5x5_2->forward(d_branch5x5_out);
 
-        float *d_branch3x3_out = branch3x3dbl_1->forward(d_input);
-        d_branch3x3_out        = branch3x3dbl_2->forward(d_branch3x3_out);
-        d_branch3x3_out        = branch3x3dbl_3->forward(d_branch3x3_out);
+        float *d_branch3x3dbl_out = branch3x3dbl_1->forward(d_input);
+        d_branch3x3dbl_out        = branch3x3dbl_2->forward(d_branch3x3dbl_out);
+        d_branch3x3dbl_out        = branch3x3dbl_3->forward(d_branch3x3dbl_out);
 
         float *d_branchPool_out = branchPool_1->forward(d_input);
         d_branchPool_out        = branchPool_2->forward(d_branchPool_out);
 
         float *d_output = concat_1->forward(d_branch1x1_out, d_branch5x5_out);
-        d_output        = concat_2->forward(d_output, d_branch3x3_out);
+        d_output        = concat_2->forward(d_output, d_branch3x3dbl_out);
         d_output        = concat_3->forward(d_output, d_branchPool_out);
 
         return d_output;
@@ -169,7 +181,7 @@ class InceptionA : public CUDANet::Module {
     }
 
   private:
-    shape2d inputSize;
+    shape2d inputShape;
     int     inputChannels;
     int     poolFeatures;
 
@@ -193,21 +205,24 @@ class InceptionA : public CUDANet::Module {
 class InceptionB : public CUDANet::Module {
   public:
     InceptionB(
-        const shape2d      inputSize,
+        const shape2d      inputShape,
         const int          inputChannels,
         const std::string &prefix
     )
-        : inputSize(inputSize), inputChannels(inputChannels) {
+        : inputShape(inputShape), inputChannels(inputChannels) {
+        
+        inputSize = inputShape.first * inputShape.second * inputChannels;
+
         // Branch 3x3
         branch3x3 = new BasicConv2d(
-            inputSize, inputChannels, 384, {3, 3}, {2, 2}, {0, 0},
+            inputShape, inputChannels, 384, {3, 3}, {2, 2}, {0, 0},
             prefix + ".branch1x1"
         );
         addLayer("", branch3x3);
 
         // Branch 3x3dbl
         branch3x3dbl_1 = new BasicConv2d(
-            inputSize, inputChannels, 64, {1, 1}, {1, 1}, {0, 0},
+            inputShape, inputChannels, 64, {1, 1}, {1, 1}, {0, 0},
             prefix + ".branch3x3dbl_1"
         );
         addLayer("", branch3x3dbl_1);
@@ -223,10 +238,10 @@ class InceptionB : public CUDANet::Module {
         addLayer("", branch3x3dbl_3);
 
         branchPool = new CUDANet::Layers::MaxPooling2d(
-            inputSize, inputChannels, {3, 3}, {2, 2}, {0, 0},
+            inputShape, inputChannels, {3, 3}, {2, 2}, {0, 0},
             CUDANet::Layers::ActivationType::NONE
         );
-        addLayer(prefix + ".branchPool", branchPool);
+        addLayer(prefix + ".branch_pool", branchPool);
 
         concat_1 = new CUDANet::Layers::Concat(
             branch3x3->getOutputSize(), branch3x3dbl_3->getOutputSize()
@@ -234,6 +249,8 @@ class InceptionB : public CUDANet::Module {
         concat_2 = new CUDANet::Layers::Concat(
             concat_1->getOutputSize(), branchPool->getOutputSize()
         );
+
+        outputSize = concat_2->getOutputSize();
     }
 
     ~InceptionB() {
@@ -272,7 +289,7 @@ class InceptionB : public CUDANet::Module {
     }
 
   private:
-    shape2d inputSize;
+    shape2d inputShape;
     int     inputChannels;
 
     BasicConv2d *branch3x3;
@@ -290,22 +307,25 @@ class InceptionB : public CUDANet::Module {
 class InceptionC : public CUDANet::Module {
   public:
     InceptionC(
-        const shape2d      inputSize,
+        const shape2d      inputShape,
         const int          inputChannels,
         const int          nChannels_7x7,
         const std::string &prefix
     )
-        : inputSize(inputSize), inputChannels(inputChannels) {
+        : inputShape(inputShape), inputChannels(inputChannels) {
+        
+        inputSize = inputShape.first * inputShape.second * inputChannels;
+        
         // Branch 1x1
         branch1x1 = new BasicConv2d(
-            inputSize, inputChannels, 192, {1, 1}, {1, 1}, {0, 0},
+            inputShape, inputChannels, 192, {1, 1}, {1, 1}, {0, 0},
             prefix + ".branch1x1"
         );
         addLayer("", branch1x1);
 
         // Branch 7x7
         branch7x7_1 = new BasicConv2d(
-            inputSize, inputChannels, nChannels_7x7, {1, 1}, {1, 1}, {0, 0},
+            inputShape, inputChannels, nChannels_7x7, {1, 1}, {1, 1}, {0, 0},
             prefix + ".branch7x7_1"
         );
         addLayer("", branch7x7_1);
@@ -322,7 +342,7 @@ class InceptionC : public CUDANet::Module {
 
         // Branch 7x7dbl
         branch7x7dbl_1 = new BasicConv2d(
-            inputSize, inputChannels, nChannels_7x7, {1, 1}, {1, 1}, {0, 0},
+            inputShape, inputChannels, nChannels_7x7, {1, 1}, {1, 1}, {0, 0},
             prefix + ".branch7x7dbl_1"
         );
         addLayer("", branch7x7dbl_1);
@@ -349,13 +369,13 @@ class InceptionC : public CUDANet::Module {
 
         // Branch Pool
         branchPool_1 = new CUDANet::Layers::AvgPooling2d(
-            inputSize, inputChannels, {3, 3}, {1, 1}, {1, 1},
+            inputShape, inputChannels, {3, 3}, {1, 1}, {1, 1},
             CUDANet::Layers::ActivationType::NONE
         );
-        addLayer("", branchPool_1);
+        addLayer(prefix + ".branch_pool", branchPool_1);
         branchPool_2 = new BasicConv2d(
             branchPool_1->getOutputDims(), inputChannels, 192, {1, 1}, {1, 1},
-            {0, 0}, prefix + ".branchPool_2"
+            {0, 0}, prefix + ".branch_pool"
         );
         addLayer("", branchPool_2);
 
@@ -369,6 +389,8 @@ class InceptionC : public CUDANet::Module {
         concat_3 = new CUDANet::Layers::Concat(
             concat_2->getOutputSize(), branchPool_2->getOutputSize()
         );
+
+        outputSize = concat_3->getOutputSize();
     }
 
     ~InceptionC() {
@@ -423,7 +445,7 @@ class InceptionC : public CUDANet::Module {
     }
 
   private:
-    shape2d inputSize;
+    shape2d inputShape;
     int     inputChannels;
 
     BasicConv2d *branch1x1;
@@ -449,50 +471,53 @@ class InceptionC : public CUDANet::Module {
 class InceptionD : public CUDANet::Module {
   public:
     InceptionD(
-        const shape2d      inputSize,
+        const shape2d      inputShape,
         const int          inputChannels,
         const std::string &prefix
     )
-        : inputSize(inputSize), inputChannels(inputChannels) {
+        : inputShape(inputShape), inputChannels(inputChannels) {
+
+        inputSize = inputShape.first * inputShape.second * inputChannels;
+
         // Branch 3x3
         branch3x3_1 = new BasicConv2d(
-            inputSize, inputChannels, 192, {1, 1}, {1, 1}, {0, 0},
+            inputShape, inputChannels, 192, {1, 1}, {1, 1}, {0, 0},
             prefix + ".branch3x3"
         );
         addLayer("", branch3x3_1);
         branch3x3_2 = new BasicConv2d(
-            inputSize, 192, 320, {3, 3}, {2, 2}, {0, 0}, prefix + ".branch3x3_2"
+            inputShape, 192, 320, {3, 3}, {2, 2}, {0, 0}, prefix + ".branch3x3_2"
         );
         addLayer("", branch3x3_2);
 
         // Branch 7x7x3
         branch7x7x3_1 = new BasicConv2d(
-            inputSize, inputChannels, 192, {1, 1}, {1, 1}, {0, 0},
+            inputShape, inputChannels, 192, {1, 1}, {1, 1}, {0, 0},
             prefix + ".branch7x7x3_1"
         );
         addLayer("", branch7x7x3_1);
         branch7x7x3_2 = new BasicConv2d(
-            inputSize, 192, 192, {1, 7}, {1, 1}, {0, 3},
+            inputShape, 192, 192, {1, 7}, {1, 1}, {0, 3},
             prefix + ".branch7x7x3_2"
         );
         addLayer("", branch7x7x3_2);
         branch7x7x3_3 = new BasicConv2d(
-            inputSize, 192, 192, {7, 1}, {1, 1}, {3, 0},
+            inputShape, 192, 192, {7, 1}, {1, 1}, {3, 0},
             prefix + ".branch7x7x3_3"
         );
         addLayer("", branch7x7x3_3);
         branch7x7x3_4 = new BasicConv2d(
-            inputSize, 192, 192, {3, 3}, {2, 2}, {0, 0},
+            inputShape, 192, 192, {3, 3}, {2, 2}, {0, 0},
             prefix + ".branch7x7x3_4"
         );
         addLayer("", branch7x7x3_4);
 
         // Branch Pool
         branchPool = new CUDANet::Layers::MaxPooling2d(
-            inputSize, 192, {3, 3}, {2, 2}, {0, 0},
+            inputShape, 192, {3, 3}, {2, 2}, {0, 0},
             CUDANet::Layers::ActivationType::NONE
         );
-        addLayer("", branchPool);
+        addLayer(prefix + ".branch_pool", branchPool);
 
         // Concat
         concat_1 = new CUDANet::Layers::Concat(
@@ -501,6 +526,8 @@ class InceptionD : public CUDANet::Module {
         concat_2 = new CUDANet::Layers::Concat(
             concat_1->getOutputSize(), branchPool->getOutputSize()
         );
+
+        outputSize = concat_2->getOutputSize();
     }
 
     ~InceptionD() {
@@ -542,7 +569,7 @@ class InceptionD : public CUDANet::Module {
     }
 
   private:
-    shape2d inputSize;
+    shape2d inputShape;
     int     inputChannels;
 
     BasicConv2d *branch3x3_1;
@@ -562,72 +589,79 @@ class InceptionD : public CUDANet::Module {
 class InceptionE : public CUDANet::Module {
   public:
     InceptionE(
-        const shape2d      inputSize,
+        const shape2d      inputShape,
         const int          inputChannels,
         const std::string &prefix
     )
-        : inputSize(inputSize), inputChannels(inputChannels) {
+        : inputShape(inputShape), inputChannels(inputChannels) {
+
+        inputSize = inputShape.first * inputShape.second * inputChannels;
+
         // Branch 1x1
         branch1x1 = new BasicConv2d(
-            inputSize, inputChannels, 320, {1, 1}, {1, 1}, {0, 0},
+            inputShape, inputChannels, 320, {1, 1}, {1, 1}, {0, 0},
             prefix + ".branch1x1"
         );
         addLayer("", branch1x1);
 
         // Branch 3x3
         branch3x3_1 = new BasicConv2d(
-            inputSize, inputChannels, 384, {1, 1}, {1, 1}, {0, 0},
+            inputShape, inputChannels, 384, {1, 1}, {1, 1}, {0, 0},
             prefix + ".branch3x3_1"
         );
         addLayer("", branch3x3_1);
         branch3x3_2a = new BasicConv2d(
-            inputSize, 384, 384, {1, 3}, {1, 1}, {0, 1},
+            inputShape, 384, 384, {1, 3}, {1, 1}, {0, 1},
             prefix + ".branch3x3_2a"
         );
         addLayer("", branch3x3_2a);
         branch3x3_2b = new BasicConv2d(
-            inputSize, 384, 384, {3, 1}, {1, 1}, {1, 0},
+            inputShape, 384, 384, {3, 1}, {1, 1}, {1, 0},
             prefix + ".branch3x3_2b"
         );
         addLayer("", branch3x3_2b);
         branch_3x3_2_concat = new CUDANet::Layers::Concat(
             branch3x3_2a->getOutputSize(), branch3x3_2b->getOutputSize()
         );
+        std::cout << "branch_3x3_2_concat: "
+                  << branch_3x3_2_concat->getOutputSize() << std::endl;
 
         // Branch 3x3dbl
         branch3x3dbl_1 = new BasicConv2d(
-            inputSize, inputChannels, 448, {1, 1}, {1, 1}, {0, 0},
+            inputShape, inputChannels, 448, {1, 1}, {1, 1}, {0, 0},
             prefix + ".branch3x3dbl_1"
         );
         addLayer("", branch3x3dbl_1);
         branch3x3dbl_2 = new BasicConv2d(
-            inputSize, 448, 384, {3, 3}, {1, 1}, {1, 1},
+            inputShape, 448, 384, {3, 3}, {1, 1}, {1, 1},
             prefix + ".branch3x3dbl_2"
         );
         addLayer("", branch3x3dbl_2);
         branch3x3dbl_3a = new BasicConv2d(
-            inputSize, 384, 384, {1, 3}, {1, 1}, {0, 1},
+            inputShape, 384, 384, {1, 3}, {1, 1}, {0, 1},
             prefix + ".branch3x3dbl_3a"
         );
         addLayer("", branch3x3dbl_3a);
         branch3x3dbl_3b = new BasicConv2d(
-            inputSize, 384, 384, {3, 1}, {1, 1}, {1, 0},
+            inputShape, 384, 384, {3, 1}, {1, 1}, {1, 0},
             prefix + ".branch3x3dbl_3b"
         );
         addLayer("", branch3x3dbl_3b);
         branch_3x3dbl_3_concat = new CUDANet::Layers::Concat(
             branch3x3dbl_3a->getOutputSize(), branch3x3dbl_3b->getOutputSize()
         );
+        std::cout << "branch_3x3dbl_3_concat: "
+                  << branch_3x3dbl_3_concat->getOutputSize() << std::endl;
 
         // Branch Pool
         branchPool_1 = new CUDANet::Layers::AvgPooling2d(
-            inputSize, inputChannels, {3, 3}, {1, 1}, {1, 1},
+            inputShape, inputChannels, {3, 3}, {1, 1}, {1, 1},
             CUDANet::Layers::ActivationType::NONE
         );
-        addLayer("", branchPool_1);
+        addLayer(prefix + ".branch_pool", branchPool_1);
         branchPool_2 = new BasicConv2d(
-            inputSize, inputChannels, 192, {1, 1}, {1, 1}, {0, 0},
-            prefix + ".branchPool_2"
+            inputShape, inputChannels, 192, {1, 1}, {1, 1}, {0, 0},
+            prefix + ".branch_pool"
         );
         addLayer("", branchPool_2);
 
@@ -641,6 +675,8 @@ class InceptionE : public CUDANet::Module {
         concat_3 = new CUDANet::Layers::Concat(
             concat_2->getOutputSize(), branchPool_2->getOutputSize()
         );
+
+        outputSize = concat_3->getOutputSize();
     }
 
     ~InceptionE() {
@@ -704,7 +740,7 @@ class InceptionE : public CUDANet::Module {
     }
 
   private:
-    shape2d inputSize;
+    shape2d inputShape;
     int     inputChannels;
 
     BasicConv2d *branch1x1;
@@ -731,24 +767,24 @@ class InceptionE : public CUDANet::Module {
 class InceptionV3 : public CUDANet::Model {
   public:
     InceptionV3(
-        const shape2d inputSize,
+        const shape2d inputShape,
         const int     inputChannels,
         const int     outputSize
     )
-        : CUDANet::Model(inputSize, inputChannels, outputSize) {
+        : CUDANet::Model(inputShape, inputChannels, outputSize) {
         conv2d_1a_3x3 = new BasicConv2d(
-            inputSize, inputChannels, 32, {3, 3}, {2, 2}, {0, 0},
-            "conv2d_1a_3x3"
+            inputShape, inputChannels, 32, {3, 3}, {2, 2}, {0, 0},
+            "Conv2d_1a_3x3"
         );
         addLayer("", conv2d_1a_3x3);
         conv2d_2a_3x3 = new BasicConv2d(
             conv2d_1a_3x3->getOutputDims(), 32, 32, {3, 3}, {1, 1}, {0, 0},
-            "conv2d_2a_3x3"
+            "Conv2d_2a_3x3"
         );
         addLayer("", conv2d_2a_3x3);
         conv2d_2b_3x3 = new BasicConv2d(
             conv2d_2a_3x3->getOutputDims(), 32, 64, {3, 3}, {1, 1}, {1, 1},
-            "conv2d_2b_3x3"
+            "Conv2d_2b_3x3"
         );
         addLayer("", conv2d_2b_3x3);
 
@@ -756,16 +792,16 @@ class InceptionV3 : public CUDANet::Model {
             conv2d_2b_3x3->getOutputDims(), 64, {3, 3}, {2, 2}, {0, 0},
             CUDANet::Layers::ActivationType::NONE
         );
-        addLayer("maxpool1", maxpool1);
+        addLayer("Maxpool1", maxpool1);
 
         conv2d_3b_1x1 = new BasicConv2d(
             maxpool1->getOutputDims(), 64, 80, {1, 1}, {1, 1}, {0, 0},
-            "conv2d_3b_1x1"
+            "Conv2d_3b_1x1"
         );
         addLayer("", conv2d_3b_1x1);
         conv2d_4a_3x3 = new BasicConv2d(
             conv2d_3b_1x1->getOutputDims(), 80, 192, {3, 3}, {1, 1}, {0, 0},
-            "conv2d_4a_3x3"
+            "Conv2d_4a_3x3"
         );
         addLayer("", conv2d_4a_3x3);
 
@@ -773,7 +809,7 @@ class InceptionV3 : public CUDANet::Model {
             conv2d_4a_3x3->getOutputDims(), 192, {3, 3}, {2, 2}, {0, 0},
             CUDANet::Layers::ActivationType::NONE
         );
-        addLayer("maxpool2", maxpool2);
+        addLayer("Maxpool2", maxpool2);
 
         Mixed_5b =
             new InceptionA(maxpool2->getOutputDims(), 192, 32, "Mixed_5b");
@@ -816,9 +852,9 @@ class InceptionV3 : public CUDANet::Model {
         addLayer("fc", fc);
     }
 
-    float* predict(const float* input) {
+    float *predict(const float *input) {
         float *d_x = inputLayer->forward(input);
-        
+
         d_x = conv2d_1a_3x3->forward(d_x);
         d_x = conv2d_2a_3x3->forward(d_x);
         d_x = conv2d_2b_3x3->forward(d_x);
@@ -839,7 +875,7 @@ class InceptionV3 : public CUDANet::Model {
         d_x = Mixed_7c->forward(d_x);
         d_x = fc->forward(d_x);
 
-        float* output = outputLayer->forward(d_x);
+        float *output = outputLayer->forward(d_x);
         return output;
     }
 
@@ -865,7 +901,6 @@ class InceptionV3 : public CUDANet::Model {
         delete fc;
     }
 
-
   private:
     BasicConv2d *conv2d_1a_3x3;
     BasicConv2d *conv2d_2a_3x3;
@@ -898,12 +933,10 @@ class InceptionV3 : public CUDANet::Model {
 };
 
 int main(int argc, const char *const argv[]) {
-
     InceptionV3 *inception_v3 = new InceptionV3({299, 299}, 3, 1000);
 
     inception_v3->printSummary();
 
-
     if (argc != 3) {
         std::cerr << "Usage: " << argv[0] << "<model_weights_path> <image_path>"
                   << std::endl;
diff --git a/include/model/module.hpp b/include/model/module.hpp
index 4609b20..57c641e 100644
--- a/include/model/module.hpp
+++ b/include/model/module.hpp
@@ -17,19 +17,14 @@ class Module : public Layers::SequentialLayer {
     int getInputSize();
 
     void addLayer(const std::string& name, Layers::SequentialLayer* layer);
-    Layers::SequentialLayer* getLayer(const std::string& name);
 
-    const std::unordered_map<std::string, Layers::SequentialLayer*>& getLayers() const;
+    const std::vector<std::pair<std::string, Layers::SequentialLayer*>>& getLayers() const;
 
   protected:
-    int inputSize;
-    int inputChannels;
+    std::vector<std::pair<std::string, Layers::SequentialLayer*>> layers;
 
     int outputSize;
-    int outputChannels;
-
-    std::vector<std::pair<std::string, Layers::SequentialLayer*>> layers;
-    std::unordered_map<std::string, Layers::SequentialLayer*>     layerMap;
+    int inputSize;
 };
 
 }  // namespace CUDANet
diff --git a/src/model/model.cpp b/src/model/model.cpp
index 5386edd..89d65ac 100644
--- a/src/model/model.cpp
+++ b/src/model/model.cpp
@@ -67,7 +67,6 @@ void Model::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
         return;
     }
 
-
     layers.push_back({name, layer});
     layerMap[name] = layer;
 }
diff --git a/src/model/module.cpp b/src/model/module.cpp
index ea82e6c..92c6b06 100644
--- a/src/model/module.cpp
+++ b/src/model/module.cpp
@@ -1,5 +1,7 @@
 #include "module.hpp"
 
+#include <algorithm>
+
 #include "cuda_helper.cuh"
 
 using namespace CUDANet;
@@ -9,7 +11,6 @@ void Module::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
 
     if (module != nullptr) {
         for (const auto& moduleLayer : module->getLayers()) {
-            layerMap[moduleLayer.first] = moduleLayer.second;
             layers.push_back({moduleLayer.first, moduleLayer.second});
         }
 
@@ -17,24 +18,17 @@ void Module::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
     }
 
     layers.push_back({name, layer});
-    layerMap[name] = layer;
-
-    // std::cout << "Wat?! - module" << name << std::endl;
 }
 
-Layers::SequentialLayer* Module::getLayer(const std::string& name) {
-    return layerMap[name];
-}
-
-const std::unordered_map<std::string, Layers::SequentialLayer*>&
+const std::vector<std::pair<std::string, Layers::SequentialLayer*>>&
 Module::getLayers() const {
-    return layerMap;
+    return layers;
 }
 
 int Module::getInputSize() {
-    return inputSize * inputSize * inputChannels;
+    return inputSize;
 }
 
 int Module::getOutputSize() {
-    return outputSize * outputSize * outputChannels;
+    return outputSize;
 }
\ No newline at end of file