Wrap test

Author: dmitriy-serdyuk

tests/extensions/test_saveload.py

@@ -42,20 +42,21 @@ def test_checkpointing():
         algorithm=GradientDescent(cost=cost, parameters=[W]),
         extensions=[Load('myweirdmodel.tar')]
     )
-    main_loop.extensions[0].main_loop = main_loop
-    main_loop._run_extensions('before_training')
-    assert_allclose(W.get_value(), old_value)
+    with main_loop.log:
+        main_loop.extensions[0].main_loop = main_loop
+        main_loop._run_extensions('before_training')
+        assert_allclose(W.get_value(), old_value)
 
-    # Make sure things work too if the model was never saved before
-    main_loop = MainLoop(
-        model=Model(cost),
-        data_stream=data_stream,
-        algorithm=GradientDescent(cost=cost, parameters=[W]),
-        extensions=[Load('mynonexisting.tar')]
-    )
-    main_loop.extensions[0].main_loop = main_loop
-    main_loop._run_extensions('before_training')
+        # Make sure things work too if the model was never saved before
+        main_loop = MainLoop(
+            model=Model(cost),
+            data_stream=data_stream,
+            algorithm=GradientDescent(cost=cost, parameters=[W]),
+            extensions=[Load('mynonexisting.tar')]
+        )
+        main_loop.extensions[0].main_loop = main_loop
+        main_loop._run_extensions('before_training')
 
-    # Cleaning
-    if os.path.exists('myweirdmodel.tar'):
-        os.remove('myweirdmodel.tar')
+        # Cleaning
+        if os.path.exists('myweirdmodel.tar'):
+            os.remove('myweirdmodel.tar')