Fix UTs which were broken due to additional observability warnings (#862

) Summary: Pull Request resolved: #862 Fix UTs which were broken due to additional observability warnings. DCP test should only validate the first warning in the list of warnings. Reviewed By: galrotem Differential Revision: D59684986 fbshipit-source-id: fa0eccc2f0a85e84c518cd15b4dfa09e34bb78ad
pytorch · Jul 12, 2024 · 53c6f91 · 53c6f91
1 parent 60059ef
commit 53c6f91
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 8 deletions.
diff --git a/tests/framework/callbacks/test_dcp_saver.py b/tests/framework/callbacks/test_dcp_saver.py
@@ -134,10 +134,8 @@ def test_save_restore_dataloader_state(self) -> None:
                 # load_state_dict is not called again on dataloader because there is no dataloader in manifest
                 self.assertEqual(stateful_dataloader.load_state_dict_call_count, 1)
                 self.assertEqual(
-                    log.output,
-                    [
-                        "WARNING:torchtnt.utils.rank_zero_log:train_dataloader was passed to `restore` but no train dataloader exists in the Snapshot"
-                    ],
+                    log.output[0],
+                    "WARNING:torchtnt.utils.rank_zero_log:train_dataloader was passed to `restore` but no train dataloader exists in the Snapshot",
                 )
 
     def test_restore_from_latest(self) -> None:

diff --git a/tests/framework/callbacks/test_torchsnapshot_saver.py b/tests/framework/callbacks/test_torchsnapshot_saver.py
@@ -126,10 +126,8 @@ def test_save_restore_dataloader_state(self) -> None:
                 # load_state_dict is not called again on dataloader because there is no dataloader in manifest
                 self.assertEqual(stateful_dataloader.load_state_dict_call_count, 1)
                 self.assertEqual(
-                    log.output,
-                    [
-                        "WARNING:torchtnt.utils.rank_zero_log:train_dataloader was passed to `restore` but no train dataloader exists in the Snapshot"
-                    ],
+                    log.output[0],
+                    "WARNING:torchtnt.utils.rank_zero_log:train_dataloader was passed to `restore` but no train dataloader exists in the Snapshot",
                 )
 
     def test_restore_from_latest(self) -> None: