diff --git a/include/model/model.hpp b/include/model/model.hpp
index 406cb8e..30e2a43 100644
--- a/include/model/model.hpp
+++ b/include/model/model.hpp
@@ -19,6 +19,8 @@ class Model {
 
   private:
 
+    Layers::Input *inputLayer;
+
     int inputSize;
     int inputChannels;
 
diff --git a/src/model/model.cpp b/src/model/model.cpp
index 0291e8b..82f441f 100644
--- a/src/model/model.cpp
+++ b/src/model/model.cpp
@@ -12,19 +12,20 @@ Model::Model(const int inputSize, const int inputChannels)
 
 
     const int inputLayerSize = inputSize * inputSize * inputChannels;
-    Layers::Input* inputLayer = new Layers::Input(inputLayerSize);
-    
-    layers.push_back(inputLayer);
+    inputLayer = new Layers::Input(inputLayerSize);
 };
 
 Model::~Model(){};
 
 float* Model::predict(const float* input) {
+
+    float* d_input = inputLayer->forward(input);
     
     for (auto& layer : layers) {
-        input = layer->forward(input);
+        d_input = layer->forward(d_input);
     }
 
+    return d_input;
 }
 
 void Model::addLayer(const std::string& name, Layers::SequentialLayer* layer) {