diff --git a/python/paddle/fluid/tests/unittests/test_device_guard.py b/python/paddle/fluid/tests/unittests/test_device_guard.py index eb8861f1bc462e61d4bd685fca39452bfd8a8c4e..330065ecd92f1537e10df7b993682da821a46548 100644 --- a/python/paddle/fluid/tests/unittests/test_device_guard.py +++ b/python/paddle/fluid/tests/unittests/test_device_guard.py @@ -33,6 +33,14 @@ def execute(main_program, startup_program): exe.run(main_program) +def get_vaild_warning_num(warning, w): + num = 0 + for i in range(len(w)): + if warning in str(w[i].message): + num += 1 + return num + + class TestDeviceGuard(unittest.TestCase): def test_device_guard(self): main_program = fluid.Program() @@ -133,7 +141,10 @@ class TestDeviceGuard(unittest.TestCase): i = fluid.layers.increment(x=i, value=1, in_place=True) fluid.layers.less_than(x=i, y=loop_len, cond=cond) - assert len(w) == 1 + warning = "The Op(while) is not support to set device." + warning_num = get_vaild_warning_num(warning, w) + assert warning_num == 1 + all_ops = main_program.global_block().ops device_attr_name = core.op_proto_and_checker_maker.kOpDeviceAttrName() for op in all_ops: @@ -169,7 +180,10 @@ class TestDeviceGuard(unittest.TestCase): shape=[1], value=4.0, dtype='float32') result = fluid.layers.less_than(x=x, y=y, force_cpu=False) - assert len(w) == 2 + warning = "\'device_guard\' has higher priority when they are used at the same time." + warning_num = get_vaild_warning_num(warning, w) + assert warning_num == 2 + all_ops = main_program.global_block().ops device_attr_name = core.op_proto_and_checker_maker.kOpDeviceAttrName() for op in all_ops: diff --git a/python/paddle/fluid/tests/unittests/test_gather_op.py b/python/paddle/fluid/tests/unittests/test_gather_op.py index 1f6e522d2668b5dcd2075ff7af6b4b1ee674632d..5dcce88acf16b96fc26cd56c2e5e034d19405b6c 100644 --- a/python/paddle/fluid/tests/unittests/test_gather_op.py +++ b/python/paddle/fluid/tests/unittests/test_gather_op.py @@ -216,7 +216,7 @@ class API_TestGather(unittest.TestCase): "index": index_np, 'axis': axis_np}, fetch_list=[out]) - expected_output = gather_numpy(x_np, index_np, axis_np) + expected_output = gather_numpy(x_np, index_np, axis_np[0]) self.assertTrue(np.allclose(result, expected_output)) diff --git a/python/paddle/fluid/tests/unittests/test_save_model_without_var.py b/python/paddle/fluid/tests/unittests/test_save_model_without_var.py index b74a6e10917f71b669a2d2e6906d6c7310c59c7e..4c63dced83b199acaffc3e703785b3eb8ec8913b 100644 --- a/python/paddle/fluid/tests/unittests/test_save_model_without_var.py +++ b/python/paddle/fluid/tests/unittests/test_save_model_without_var.py @@ -50,7 +50,7 @@ class TestSaveModelWithoutVar(unittest.TestCase): params_filename='params') expected_warn = "no variable in your model, please ensure there are any variables in your model to save" self.assertTrue(len(w) > 0) - self.assertTrue(expected_warn == str(w[0].message)) + self.assertTrue(expected_warn == str(w[-1].message)) if __name__ == '__main__':