moved wrapper class from header file to locally defined in export fun…

…ction
microsoft · Jul 25, 2024 · e66740d · e66740d
1 parent aa97469
commit e66740d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 13 deletions.
diff --git a/orttraining/orttraining/training_api/module.cc b/orttraining/orttraining/training_api/module.cc
@@ -671,6 +671,19 @@ Status Module::ExportModelForInferencing(const std::string& inference_model_path
                 "Cannot export the model with a nominal state. Please load the model parameters first.");
   ORT_RETURN_IF(!eval_sess_, "Eval model was not provided. Cannot export a model for inferencing.");
 
+  class EvalSessionWrapper : public InferenceSession {
+   public:
+    using InferenceSession::InferenceSession;
+
+    Graph& GetMutableGraph() const {
+      return model_->MainGraph();
+    }
+
+    Model& GetMutableModel() {
+      return *model_;
+    }
+  };
+
   // Once finished_training is set to true, will no longer be able to train or evaluate with this module
   // since the eval session graph will have been modified.
   finished_training_ = true;

diff --git a/orttraining/orttraining/training_api/module.h b/orttraining/orttraining/training_api/module.h
@@ -12,19 +12,6 @@ namespace onnxruntime {
 namespace training {
 namespace api {
 
-class EvalSessionWrapper : public InferenceSession {
- public:
-  using InferenceSession::InferenceSession;
-
-  Graph& GetMutableGraph() const {
-    return model_->MainGraph();
-  }
-
-  Model& GetMutableModel() {
-    return *model_;
-  }
-};
-
 struct Parameter {
  public:
   Parameter(const std::string& name, const OrtValue& data, const bool requires_grad)