fix test

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

@@ -346,6 +346,7 @@ def get_post_processed_model(
             True,
             self._device,
             self._export_mode,
+            self._logger,
             self._export_extra_kwargs,
         )
 
@@ -389,7 +390,7 @@ def get_post_processed_model(
             flatten_inputs = []
 
             # This looks a bit duplicated with `extract_data_and_schema` function, but this might be better to
-            # defined as a specialized logic which is the counter-part of `parse_inputs_for_onnx_export`, which handles
+            # defined as a specialized logic that is the counter-part of `parse_inputs_for_onnx_export`, which handles
             # args and kwargs separately.
             for name, data_accessor in cur_model_info_for_export.onnx_graph_input_data_accessor.items():
                 d = data_accessor(copied_args, copied_kwargs)

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

@@ -6,7 +6,6 @@
 import copy
 import gc
 import inspect
-import warnings
 from collections import OrderedDict, abc
 from functools import partial
 from logging import Logger
@@ -244,6 +243,7 @@ def parse_inputs_for_onnx_export(
     constant_as_tensor: bool,
     device: torch.device,
     export_mode: int,
+    logger: Logger,
     export_extra_kwargs: Optional[Dict[str, any]] = None,
 ) -> ModelInfoForExport:
     """Parses through the model inputs and returns _InputInfo.
@@ -282,7 +282,7 @@ def _add_dynamic_shape(name, input) -> Dict[str, Dict[int, str]]:
         return dynamic_axes
 
     def _warn_of_constant_inputs(data):
-        warnings.warn(f"Received input of type {type(data)} is treated as a constant by ORT by default.")
+        logger.info(f"Received input of type {type(data)} is treated as a constant by ORT by default.")
 
     def _add_input(
         name: str, input_value, onnx_graph_input_names: List[str], cur_func: Callable, tensor_idx: List[int]

orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py

@@ -4239,14 +4239,23 @@ def test_hf_save_pretrained():
             assert p1.data.ne(p2.data).sum() == 0
 
 
-def test_ortmodule_string_inputs_are_ignored():
+def test_ortmodule_string_inputs_are_ignored(caplog):
     pt_model = MyStrNet()
-    target_str = "Received input of type <class 'str'> which may be treated as a constant by ORT by default."
-    with pytest.warns(UserWarning, match=target_str):
-        ort_model = ORTModule(copy.deepcopy(pt_model), DebugOptions(log_level=LogLevel.INFO))
-        x = torch.randn(1, 2)
-        out = ort_model(x, "hello")
-        _test_helpers.assert_values_are_close(out, x + 1)
+    target_str = "Received input of type <class 'str'> is treated as a constant by ORT by default."
+
+    ort_model = ORTModule(copy.deepcopy(pt_model), DebugOptions(log_level=LogLevel.INFO))
+    x = torch.randn(1, 2)
+    out = ort_model(x, "hello")
+    _test_helpers.assert_values_are_close(out, x + 1)
+
+    found_log = False
+    for record in caplog.records:
+        msg = record.getMessage()
+        if target_str in msg:
+            found_log = True
+            break
+
+    assert found_log, f"Expected to find log message '{target_str}' in the logs, but didn't find it."
 
 
 def test_ortmodule_list_input():