From 67fcb0c923c99bf262f91f523a18f6b1319dcce9 Mon Sep 17 00:00:00 2001 From: Youwei Song Date: Mon, 14 Oct 2019 16:43:48 +0800 Subject: [PATCH] fix api fluid.cuda_pinned_places (#20424) * fix api fluid.cuda_pinned_places, test=develop * fix api fluid.cuda_pinned_places, 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 740e86856ad..8d46fc9c113 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