Skip to content

Commit

Permalink
update getNvpPnnxModel func comments
Browse files Browse the repository at this point in the history
  • Loading branch information
sen.li committed Jun 19, 2024
1 parent 8c47d8d commit d85c333
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 7 deletions.
6 changes: 3 additions & 3 deletions tools/pnnx/src/parse/pnnx_graph_parse.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -65,9 +65,9 @@ bool PnnxGraph::getNvpPnnxModel(const std::string& pt_path, const std::string& s
std::strcpy( argv[argc - 2], end_nodes_info.c_str());

//insert extract_model_name
std::string custom_op_py_info = "extract_model_name=" + custom_op_py;
argv[argc - 1] = new char[custom_op_py_info.size() + 1];
std::strcpy(argv[argc - 2], custom_op_py_info.c_str());
std::string extract_model_name_info = "extract_model_name=" + extract_model_name;
argv[argc - 1] = new char[extract_model_name_info.size() + 1];
std::strcpy(argv[argc - 1], extract_model_name_info.c_str());

int result = main(argc, argv);

Expand Down
11 changes: 7 additions & 4 deletions tools/pnnx/src/parse/pnnx_graph_parse.h
Original file line number Diff line number Diff line change
Expand Up @@ -15,21 +15,24 @@ class PnnxGraph
* @brief Get the Nvp Pnnx Model object
*
* @param pt_path torchscript path
* @param save_dir the dir of save pnnx
* @param input_shape input shape of tensor
* @param custom_op_path the path of define custom op
* @param custom_op_py the py path of define custom op
* @param start_nodes the list of start nodes
* @param end_nodes the list of end nodes
* @param extract_model_name the name of which to extract
* @return true
* @return false
*/
bool getNvpPnnxModel(const std::string& pt_path, \
const std::string& save_dir, \
const std::string& save_dir, \
const std::string& input_shape, \
const std::string& custom_op_path, \
const std::string& custom_op_py,
const std::string& start_nodes = "",
const std::string& end_nodes = "");
const std::string& custom_op_py,\
const std::string& start_nodes = "",\
const std::string& end_nodes = "",\
const std::string& extract_model_name = "model");



Expand Down

0 comments on commit d85c333

Please sign in to comment.