Suppress stdout in test_train

tests/test_train.py

@@ -5,6 +5,8 @@
 # the root directory of this source tree. An additional grant of patent rights
 # can be found in the PATENTS file in the same directory.
 
+import contextlib
+from io import StringIO
 import unittest
 
 from unittest.mock import MagicMock, patch
@@ -37,6 +39,7 @@ class TestLoadCheckpoint(unittest.TestCase):
         [p.start() for p in self.applied_patches]
 
     def test_load_partial_checkpoint(self):
+        with contextlib.redirect_stdout(StringIO()):
             trainer = mock_trainer(2, 200, False)
             loader = mock_loader(150)
             epoch, ds = train.load_checkpoint(MagicMock(), trainer, loader)
@@ -44,6 +47,7 @@ class TestLoadCheckpoint(unittest.TestCase):
             self.assertEqual(next(ds), 50)
 
     def test_load_full_checkpoint(self):
+        with contextlib.redirect_stdout(StringIO()):
             trainer = mock_trainer(2, 300, True)
             loader = mock_loader(150)
             epoch, ds = train.load_checkpoint(MagicMock(), trainer, loader)
@@ -51,6 +55,7 @@ class TestLoadCheckpoint(unittest.TestCase):
             self.assertEqual(next(iter(ds)), 0)
 
     def test_load_no_checkpoint(self):
+        with contextlib.redirect_stdout(StringIO()):
             trainer = mock_trainer(0, 0, False)
             loader = mock_loader(150)
             self.patches['os.path.isfile'].return_value = False