From 5a64b5ade7de005ba2965835524371cd4182205b Mon Sep 17 00:00:00 2001 From: Youwei Song Date: Wed, 16 Oct 2019 18:57:40 +0800 Subject: [PATCH] cherry-pick #20424 (#20661) test=release/1.6, test=develop --- python/paddle/fluid/framework.py | 4 ++-- python/paddle/fluid/tests/unittests/test_tensor.py | 4 ++++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/python/paddle/fluid/framework.py b/python/paddle/fluid/framework.py index 610dfca8737..4191ea6ea0f 100644 --- a/python/paddle/fluid/framework.py +++ b/python/paddle/fluid/framework.py @@ -373,8 +373,8 @@ def cuda_pinned_places(device_count=None): assert core.is_compiled_with_cuda(), \ "Not compiled with CUDA" if device_count is None: - device_count = _cpu_num() - return [core.cuda_pinned_places()] * device_count + device_count = len(_cuda_ids()) + return [core.CUDAPinnedPlace()] * device_count class NameScope(object): diff --git a/python/paddle/fluid/tests/unittests/test_tensor.py b/python/paddle/fluid/tests/unittests/test_tensor.py index ec180456acf..09417a3e1e0 100644 --- a/python/paddle/fluid/tests/unittests/test_tensor.py +++ b/python/paddle/fluid/tests/unittests/test_tensor.py @@ -275,6 +275,10 @@ class TestTensor(unittest.TestCase): self.assertTrue( isinstance( tensor._mutable_data(place, dtype), numbers.Integral)) + places = fluid.cuda_pinned_places() + self.assertTrue( + isinstance( + tensor._mutable_data(places[0], dtype), numbers.Integral)) if __name__ == '__main__': -- GitLab