From b5fb205df84fe033d956672fa5e1b0710eca92bf Mon Sep 17 00:00:00 2001
From: LordMathis <matus@namesny.com>
Date: Sun, 19 May 2024 20:22:08 +0200
Subject: [PATCH] Update module add layer

---
 src/model/module.cpp | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/model/module.cpp b/src/model/module.cpp
index 7cfde81..4fd2728 100644
--- a/src/model/module.cpp
+++ b/src/model/module.cpp
@@ -5,8 +5,16 @@
 using namespace CUDANet;
 
 void Module::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
-    layers.push_back({ name, layer });
-    layerMap[name] = layer;
+    Module* module = dynamic_cast<Module*>(layer);
+
+    if (module != nullptr) {
+        layers.push_back({ name, module });
+        for (const auto& moduleLayer : module->getLayers()) {
+            layerMap[moduleLayer.first] = moduleLayer.second;
+        }
+
+        return;
+    }
 }
 
 Layers::SequentialLayer* Module::getLayer(const std::string& name) {