[Application] yolo v2 bug fix

register a custom layer that had been omitted on yolo v2. **Self evaluation:** 1. Build test: [X]Passed [ ]Failed [ ]Skipped 2. Run test: [X]Passed [ ]Failed [ ]Skipped Signed-off-by: Seungbaek Hong <[email protected]>
nnstreamer · Jun 3, 2024 · d50487a · d50487a
1 parent 2a66923
commit d50487a
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/Applications/YOLOv2/jni/main.cpp b/Applications/YOLOv2/jni/main.cpp
@@ -284,6 +284,17 @@ int main(int argc, char *argv[]) {
   std::cout << "batch_size: " << BATCH_SIZE << " epochs: " << EPOCHS
             << std::endl;
 
+  try {
+    auto &app_context = nntrainer::AppContext::Global();
+    app_context.registerFactory(nntrainer::createLayer<custom::ReorgLayer>);
+    app_context.registerFactory(
+      nntrainer::createLayer<custom::YoloV2LossLayer>);
+  } catch (std::invalid_argument &e) {
+    std::cerr << "failed to register factory, reason: " << e.what()
+              << std::endl;
+    return 1;
+  }
+
   try {
     // create YOLO v2 model
     ModelHandle model = YOLO();