[tools/onert_train] Support checkpoint loader

ONE-DCO-1.0-Signed-off-by: Jiyoung Yun <[email protected]>
Samsung · Aug 9, 2024 · bbc757d · bbc757d
1 parent 22230e9
commit bbc757d
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/tests/tools/onert_train/src/onert_train.cc b/tests/tools/onert_train/src/onert_train.cc
@@ -161,7 +161,12 @@ int main(const int argc, char **argv)
     // prepare execution
 
     // TODO When nnfw_{prepare|run} are failed, can't catch the time
-    measure.run(PhaseType::PREPARE, [&]() { NNPR_ENSURE_STATUS(nnfw_train_prepare(session)); });
+    measure.run(PhaseType::PREPARE, [&]() {
+      NNPR_ENSURE_STATUS(nnfw_train_prepare(session));
+
+      if (auto name = args.getCheckpointFilename(); name != "")
+        NNPR_ENSURE_STATUS(nnfw_train_import_checkpoint(session, name.c_str()));
+    });
 
     // prepare input and expected tensor info lists
     std::vector<nnfw_tensorinfo> input_infos;