Using ORTModule to wrap a evaluation model should not change the mode (#12747)

Using ORTModule to wrap a evaluation model should not change the mode of model

orttraining/orttraining/python/training/ortmodule/ortmodule.py

@@ -98,6 +98,7 @@ class ORTModule(torch.nn.Module):
                 override_policy=_FallbackPolicy.FALLBACK_FORCE_TORCH_FORWARD,
             )
 
+        self.train(module.training)
         # Finally, ORTModule initialization is complete.
         # Assign self._is_initialized to True after all the ORTModule class attributes have been assigned
         # else, they will be assigned to self._torch_module.original_module instead.

orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py

@@ -390,6 +390,7 @@ def run_before_tests():
 
 def _get_bert_for_sequence_classification_model(
     device,
+    is_training=False,
     output_attentions=False,
     output_hidden_states=False,
     return_dict=True,
@@ -414,6 +415,10 @@ def _get_bert_for_sequence_classification_model(
         config=config,
     ).to(device)
 
+    if is_training:
+        model.train()
+    else:
+        model.eval()
     return model
 
 
@@ -2012,7 +2017,7 @@ def test_bert_inputs_with_dynamic_shape():
 
     # create pytorch model with dropout disabled
     pt_model = _get_bert_for_sequence_classification_model(
-        "cuda", hidden_dropout_prob=0.0, attention_probs_dropout_prob=0.0
+        "cuda", is_training=True, hidden_dropout_prob=0.0, attention_probs_dropout_prob=0.0
     )
     ort_model = ORTModule(copy.deepcopy(pt_model))
 
@@ -5334,3 +5339,21 @@ def test_squeeze_custom_symbolic_registry():
     _test_helpers.assert_values_are_close(pt_prediction, ort_prediction)
     _test_helpers.assert_values_are_close(pt_loss, ort_loss)
     _test_helpers.assert_values_are_close(pt_x.grad, ort_x.grad)
+
+
+def test_eval_model_mode():
+    device = "cuda"
+    n, d_in, h_size, d_out = 64, 2, 2, 2
+    origin_model = NeuralNetSinglePositionalArgument(d_in, h_size, d_out).to(device)
+    x = torch.randn(n, d_in, device=device)
+    ort_model = ORTModule(origin_model)
+    for initial_mode in (True, False):
+        model = copy.deepcopy(ort_model)
+        model.train(initial_mode)
+        for step in range(10):
+            for new_mode in (True, False):
+                model.train(new_mode)
+                model(x)
+                assert model.training == new_mode
+                assert model._torch_module.is_training() == new_mode
+                assert model._torch_module._flattened_module.training == new_mode