Merge pull request huggingface#21 from microsoft/aibhanda/dev_options

Add ORT internal options config flag
jambayk · Aug 26, 2021 · e4e57f0 · e4e57f0
2 parents 2faedf4 + 9109e4f
commit e4e57f0
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/src/transformers/trainer.py b/src/transformers/trainer.py
@@ -1107,6 +1107,9 @@ def train(
             from torch_ort import ORTModule
             logger.info("Converting to ORTModule ....")
             model = ORTModule(self.model)
+            if args.ort_config_file:
+                from torch_ort.experimental.json_config import load_from_json
+                load_from_json(model, path=args.ort_config_file)
             self.model_wrapped = model
         if args.deepspeed:
             if args.ort:

diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py
@@ -526,6 +526,10 @@ class TrainingArguments:
         default=False,
         metadata={"help": "Enable Ort"},
     )
+    ort_config_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "Configure ORTModule internal options"},
+    )
     label_smoothing_factor: float = field(
         default=0.0, metadata={"help": "The label smoothing epsilon to apply (zero means no label smoothing)."}
     )