Skip to content

Commit

Permalink
Add convenience load methond for forward. (#5446)
Browse files Browse the repository at this point in the history
Summary:
Pull Request resolved: #5446

People write `module->load_method("forward")` too often, let's simplify that a bit.

Reviewed By: kirklandsign

Differential Revision: D62906055

fbshipit-source-id: d45934a27c61fd0ea644e6b58a21116dbc02fa17
(cherry picked from commit b14dea8)
  • Loading branch information
shoumikhin authored and pytorchbot committed Sep 18, 2024
1 parent 8a7cb9c commit fd20b22
Showing 1 changed file with 17 additions and 1 deletion.
18 changes: 17 additions & 1 deletion extension/module/module.h
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ class Module {
*
* @param[in] file_path The path to the ExecuTorch program file to load.
* @param[in] load_mode The loading mode to use.
* @param[in] event_tracer A EventTracer used for tracking and logging events.
*/
explicit Module(
const std::string& file_path,
Expand Down Expand Up @@ -132,13 +133,28 @@ class Module {
* needed. The loaded method is cached to reuse the next time it's executed.
*
* @param[in] method_name The name of the method to load.
* @param[in] event_tracer A EventTracer used for tracking and logging events.
*
* @returns An Error to indicate success or failure.
*/
ET_NODISCARD
runtime::Error load_method(
const std::string& method_name,
torch::executor::EventTracer* tracer = nullptr);
torch::executor::EventTracer* event_tracer = nullptr);

/**
* Load the 'forward' method from the program and set up memory management if
* needed. The loaded method is cached to reuse the next time it's executed.
*
* @param[in] event_tracer An event tracer used for tracking and logging
* events.
*
* @returns An Error to indicate success or failure.
*/
ET_NODISCARD inline runtime::Error load_forward(
torch::executor::EventTracer* event_tracer = nullptr) {
return load_method("forward", event_tracer);
}

/**
* Checks if a specific method is loaded.
Expand Down

0 comments on commit fd20b22

Please sign in to comment.