Set gradient as output only for easy mode (#3694)

This commit is contained in:
ashbhandare 2020-04-24 15:28:28 -07:00 committed by GitHub
parent b4d4ea2e5f
commit d06763ac1c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -94,7 +94,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 = true;
config.set_gradients_as_graph_outputs = false;
config.gradient_accumulation_steps = parameters.gradient_accumulation_steps;
@ -115,6 +115,7 @@ 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;
@ -276,4 +277,4 @@ void addObjectMethodsForTraining(py::module& m) {
}
} // namespace python
} // namespace onnxruntime
} // namespace onnxruntime