From f14f0e19eceaa32ea179829c5314676df8011a51 Mon Sep 17 00:00:00 2001
From: ashbhandare <aibhanda@microsoft.com>
Date: Wed, 22 Jun 2022 08:48:19 -0700
Subject: [PATCH] Fix trainer (#11933)

fix trainer:
---
 .../orttraining/test/training_api/trainer/trainer.cc        | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/orttraining/orttraining/test/training_api/trainer/trainer.cc b/orttraining/orttraining/test/training_api/trainer/trainer.cc
index d9eac9dfe9..6d6d8ade32 100644
--- a/orttraining/orttraining/test/training_api/trainer/trainer.cc
+++ b/orttraining/orttraining/test/training_api/trainer/trainer.cc
@@ -25,6 +25,7 @@ struct TestRunnerParameters {
   std::string optimizer_training_graph_path;
   std::string checkpoint_to_load_path;
   std::string model_name;
+  std::string synthetic_input_type;
 
   // Data configs.
   std::string train_data_dir;
@@ -74,6 +75,8 @@ bool ParseArguments(int argc, char* argv[], TestRunnerParameters& params) {
         cxxopts::value<std::string>()->default_value(""))
       ("model_name", "The name of the model.",
         cxxopts::value<std::string>()->default_value("model_test"))
+      ("synthetic_input_type", "Input type can be 'dummy'(input, target) or 'attention'(input_id, attention, target)",
+        cxxopts::value<std::string>()->default_value("attention"))
 
       ("train_data_dir", "Input ONNX example files (can be a glob or comma separated).",
         cxxopts::value<std::string>()->default_value("bert_data/128/books_wiki_en_corpus/train"))
@@ -102,6 +105,7 @@ bool ParseArguments(int argc, char* argv[], TestRunnerParameters& params) {
     } else {
       params.model_evaluation_graph_path = eval_path;
     }
+    params.synthetic_input_type = flags["synthetic_input_type"].as<std::string>();
 
     params.optimizer_training_graph_path = flags["optimizer_training_graph_path"].as<std::string>();
     params.checkpoint_to_load_path = flags["checkpoint_to_load_path"].as<std::string>();
@@ -142,7 +146,7 @@ void InitSyntheticDataLoader(
     onnxruntime::training::test::training_api::SyntheticDataLoader& data_loader,
     const TestRunnerParameters& params,
     int64_t num_of_batches_per_epoch) {
-  bool sample_model = false;
+  bool sample_model = params.synthetic_input_type == "dummy" ? true : false;
   if (sample_model) {
     std::vector<int64_t> input1_shape{params.train_batch_size, 784};
     std::vector<int64_t> target_shape{params.train_batch_size};