diff --git a/include/model/module.hpp b/include/model/module.hpp
deleted file mode 100644
index e6ae17b..0000000
--- a/include/model/module.hpp
+++ /dev/null
@@ -1,32 +0,0 @@
-#ifndef CUDANET_MODULE_H
-#define CUDANET_MODULE_H
-
-#include <string>
-#include <unordered_map>
-#include <vector>
-
-#include "layer.hpp"
-
-namespace CUDANet {
-
-class Module : public Layers::SequentialLayer {
-  public:
-    virtual float* forward(const float* d_input) = 0;
-
-    int getOutputSize();
-    int getInputSize();
-
-    void addLayer(const std::string& name, Layers::SequentialLayer* layer);
-
-    const std::vector<std::pair<std::string, Layers::SequentialLayer*>>& getLayers() const;
-
-  protected:
-    std::vector<std::pair<std::string, Layers::SequentialLayer*>> layers;
-
-    int outputSize;
-    int inputSize;
-};
-
-}  // namespace CUDANet
-
-#endif
\ No newline at end of file
diff --git a/include/module.hpp b/include/module.hpp
new file mode 100644
index 0000000..4521741
--- /dev/null
+++ b/include/module.hpp
@@ -0,0 +1,35 @@
+#pragma once
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "layer.hpp"
+
+namespace CUDANet {
+
+class Module {
+  public:
+    CUDANet::Shape input_shape();
+
+    CUDANet::Shape output_shape();
+
+    size_t input_size();
+
+    size_t output_size();
+
+    void register_layer(const std::string& name, Layer& layer);
+
+    void register_module(Module& module);
+
+    const std::vector<std::pair<std::string, Layer&>>& get_layers() const;
+
+  protected:
+    std::vector<std::pair<std::string, Layer&>> layers;
+
+    CUDANet::Shape in_shape;
+    CUDANet::Shape out_shape;
+};
+
+}  // namespace CUDANet
diff --git a/src/module.cpp b/src/module.cpp
index 1fb5434..3ea9e00 100644
--- a/src/module.cpp
+++ b/src/module.cpp
@@ -4,29 +4,41 @@
 
 using namespace CUDANet;
 
-void Module::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
-    const Module* module = dynamic_cast<Module*>(layer);
+CUDANet::Shape Module::input_shape() {
+    return in_shape;
+}
 
-    if (module != nullptr) {
-        for (const auto& moduleLayer : module->getLayers()) {
-            layers.push_back({moduleLayer.first, moduleLayer.second});
-        }
+CUDANet::Shape Module::output_shape() {
+    return out_shape;
+}
 
-        return;
+size_t Module::input_size() {
+    size_t count = 1;
+    for (const auto& dim : in_shape) {
+        count *= dim;
     }
+    return sizeof(float) * count;
+}
 
+size_t Module::output_size() {
+    size_t count = 1;
+    for (const auto& dim : out_shape) {
+        count *= dim;
+    }
+    return sizeof(float) * count;
+}
+
+void Module::register_layer(const std::string& name, Layer& layer) {
     layers.push_back({name, layer});
 }
 
-const std::vector<std::pair<std::string, Layers::SequentialLayer*>>&
-Module::getLayers() const {
+void Module::register_module(Module& module) {
+    for (const auto& moduleLayer : module.get_layers()) {
+        layers.push_back({moduleLayer.first, moduleLayer.second});
+    }
+}
+
+const std::vector<std::pair<std::string, Layer&>>&
+Module::get_layers() const {
     return layers;
 }
-
-int Module::getInputSize() {
-    return inputSize;
-}
-
-int Module::getOutputSize() {
-    return outputSize;
-}
\ No newline at end of file