Fix adding module layers

This commit is contained in:
2024-05-29 21:02:26 +02:00
parent 4159454842
commit 8ac2da004c
4 changed files with 31 additions and 20 deletions

View File

@@ -56,17 +56,18 @@ float* Model::predict(const float* input) {
}
void Model::addLayer(const std::string& name, Layers::SequentialLayer* layer) {
Module* module = dynamic_cast<Module*>(layer);
const 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;
layers.push_back({moduleLayer.first, moduleLayer.second});
}
return;
}
layers.push_back({name, layer});
layerMap[name] = layer;
}