diff --git a/python/paddle/fluid/framework.py b/python/paddle/fluid/framework.py index 740e86856ad632d553c6f7ee9854aa6b7689207a..8d46fc9c1136ab9aa0b90c8f8915fb4e900a44dd 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 ec180456acf00c70a9593ea12b728fe1b65335c1..09417a3e1e0752ca2104fb5032ade2b634c863c2 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__':