diff --git a/python/paddle/fluid/tests/unittests/test_imperative_numpy_bridge.py b/python/paddle/fluid/tests/unittests/test_imperative_numpy_bridge.py index 4f3089baffdd3fd57c069a178205fced6044a448..7b8d31ff030e503f872b9afd923ce4c6252a026a 100644 --- a/python/paddle/fluid/tests/unittests/test_imperative_numpy_bridge.py +++ b/python/paddle/fluid/tests/unittests/test_imperative_numpy_bridge.py @@ -42,6 +42,7 @@ class TestImperativeNumpyBridge(unittest.TestCase): self.assertEqual(data_np[0][0], -1) if _in_eager_mode(): # eager_mode, var2 is EagerTensor, is not subscriptable + # TODO(wuweilong): to support slice in eager mode later self.assertNotEqual(var2.numpy()[0][0], -1) else: self.assertNotEqual(var2[0][0].numpy()[0], -1) diff --git a/python/paddle/fluid/tests/unittests/test_imperative_using_non_zero_gpu.py b/python/paddle/fluid/tests/unittests/test_imperative_using_non_zero_gpu.py index f2dfaef397797a9d1b7e2900cdb2cc3617bcb933..46a89efcec491e7a6938a1d046fc8bb0ba8a7b91 100644 --- a/python/paddle/fluid/tests/unittests/test_imperative_using_non_zero_gpu.py +++ b/python/paddle/fluid/tests/unittests/test_imperative_using_non_zero_gpu.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import paddle import paddle.fluid as fluid import unittest -from paddle.fluid.dygraph import to_variable, Embedding, guard +from paddle.fluid.dygraph import to_variable, guard import numpy as np +from paddle.fluid.framework import _test_eager_guard class TestImperativeUsingNonZeroGpu(unittest.TestCase): @@ -24,12 +26,21 @@ class TestImperativeUsingNonZeroGpu(unittest.TestCase): var = to_variable(np_arr) self.assertTrue(np.array_equal(np_arr, var.numpy())) - def test_non_zero_gpu(self): + def func_non_zero_gpu(self): if not fluid.is_compiled_with_cuda(): return np_arr = np.random.random([11, 13]).astype('float32') - self.run_main(np_arr, fluid.CUDAPlace(0)) + if paddle.device.cuda.device_count() > 1: + # should use non zero gpu if there are more than 1 gpu + self.run_main(np_arr, fluid.CUDAPlace(1)) + else: + self.run_main(np_arr, fluid.CUDAPlace(0)) + + def test_non_zero_gpu(self): + with _test_eager_guard(): + self.func_non_zero_gpu() + self.func_non_zero_gpu() if __name__ == '__main__':