export weights as a constants in graph, so can do constant folding to…

… them
microsoft · Jan 26, 2024 · b70cd75 · b70cd75
1 parent 656ca66
commit b70cd75
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py
@@ -402,7 +402,7 @@ def _get_exported_model(self, input_schema: ORTModelInputOutputSchemaType, *inpu
                     "training": self._export_mode,
                     "dynamic_axes": self._input_info.dynamic_axes,
                     "verbose": torch_exporter_verbose_log,
-                    "export_params": False,
+                    "export_params": False if self._export_mode != torch.onnx.TrainingMode.EVAL else True,
                     "keep_initializers_as_inputs": True,
                 }