[tools/onert_train] Use optimizer type argument (#11171)

This commit uses optimizer type argument and passes it to nnfw_train_info structure. ONE-DCO-1.0-Signed-off-by: Jiyoung Yun <[email protected]>
Samsung · Jul 27, 2023 · e476e73 · e476e73
1 parent f66855e
commit e476e73
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/tests/tools/onert_train/src/onert_train.cc b/tests/tools/onert_train/src/onert_train.cc
@@ -129,11 +129,25 @@ int main(const int argc, char **argv)
       }
     };
 
+    auto convertOptType = [](int type) {
+      switch (type)
+      {
+        case 0:
+          return NNFW_TRAIN_OPTIMIZER_SGD;
+        case 1:
+          return NNFW_TRAIN_OPTIMIZER_ADAM;
+        default:
+          std::cerr << "E: not supported optimizer type" << std::endl;
+          exit(-1);
+      }
+    };
+
     // prepare training info
     nnfw_train_info tri;
     tri.batch_size = args.getBatchSize();
     tri.learning_rate = args.getLearningRate();
     tri.loss = convertLossType(args.getLossType());
+    tri.opt = convertOptType(args.getOptimizerType());
 
     // prepare execution