From 424a00bf042b96b864b936324df69c4f2695dee4 Mon Sep 17 00:00:00 2001
From: ashbhandare <aibhanda@microsoft.com>
Date: Thu, 7 May 2020 15:07:14 -0700
Subject: [PATCH] Fix enabling gradient as output for easy mode. (#3866)

---
 orttraining/orttraining/python/ort_trainer.py              | 2 ++
 orttraining/orttraining/python/orttraining_pybind_state.cc | 7 ++++---
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/orttraining/orttraining/python/ort_trainer.py b/orttraining/orttraining/python/ort_trainer.py
index 7e0e678026..6626d88dda 100644
--- a/orttraining/orttraining/python/ort_trainer.py
+++ b/orttraining/orttraining/python/ort_trainer.py
@@ -405,6 +405,7 @@ def create_ort_training_session_with_optimizer(model, device, training_optimizer
     ort_parameters.allreduce_post_accumulation = allreduce_post_accumulation
     ort_parameters.partition_optimizer = partition_optimizer
     ort_parameters.enable_grad_norm_clip = enable_grad_norm_clip
+    ort_parameters.set_gradients_as_graph_outputs = False
 
     output_types = {}
     for output in model.graph.output:
@@ -477,6 +478,7 @@ def create_ort_training_session_bind_parameters(model, device, world_rank=-1, wo
     ort_parameters.world_rank = world_rank
     ort_parameters.world_size = world_size
     ort_parameters.gradient_accumulation_steps = gradient_accumulation_steps
+    ort_parameters.set_gradients_as_graph_outputs = True
 
     torch_params = {}
     output_types = {}
diff --git a/orttraining/orttraining/python/orttraining_pybind_state.cc b/orttraining/orttraining/python/orttraining_pybind_state.cc
index 4628bad416..f841e0487f 100644
--- a/orttraining/orttraining/python/orttraining_pybind_state.cc
+++ b/orttraining/orttraining/python/orttraining_pybind_state.cc
@@ -51,6 +51,7 @@ struct TrainingParameters {
   int horizontal_parallel_size = 1;
   bool partition_optimizer = false;
   bool enable_grad_norm_clip = true;
+  bool set_gradients_as_graph_outputs = false;
 };
 
 struct TrainingConfigurationResult {
@@ -94,7 +95,7 @@ TrainingConfigurationResult ConfigureSessionForTraining(
   config.weight_names_to_not_train = parameters.weights_not_to_train;
   config.immutable_weights = parameters.immutable_weights;
 
-  config.set_gradients_as_graph_outputs = false;
+  config.set_gradients_as_graph_outputs = parameters.set_gradients_as_graph_outputs;
 
   config.gradient_accumulation_steps = parameters.gradient_accumulation_steps;
 
@@ -115,7 +116,6 @@ TrainingConfigurationResult ConfigureSessionForTraining(
   config.loss_name = parameters.loss_output_name;
 
   if (!parameters.training_optimizer_name.empty()) {
-    config.set_gradients_as_graph_outputs = true;
     training::TrainingSession::TrainingConfiguration::OptimizerConfiguration opt{};
     opt.name = parameters.training_optimizer_name;
     opt.learning_rate_input_name = parameters.lr_params_feed_name;
@@ -181,7 +181,8 @@ void addObjectMethodsForTraining(py::module& m) {
       .def_readwrite("world_size", &TrainingParameters::world_size)
       .def_readwrite("gradient_accumulation_steps", &TrainingParameters::gradient_accumulation_steps)
       .def_readwrite("partition_optimizer", &TrainingParameters::partition_optimizer)
-      .def_readwrite("enable_grad_norm_clip", &TrainingParameters::enable_grad_norm_clip);
+      .def_readwrite("enable_grad_norm_clip", &TrainingParameters::enable_grad_norm_clip)
+      .def_readwrite("set_gradients_as_graph_outputs", &TrainingParameters::set_gradients_as_graph_outputs);
 
   py::class_<TrainingConfigurationResult> config_result(m, "TrainingConfigurationResult", "pbdoc(Configuration result for training.)pbdoc");
   config_result.def(py::init())