Skip to content

Commit 5a64b5a

Browse files
songyouweichengduoZH
authored andcommitted
cherry-pick #20424 (#20661)
test=release/1.6, test=develop
1 parent 3880f3d commit 5a64b5a

File tree

2 files changed

+6
-2
lines changed

2 files changed

+6
-2
lines changed

python/paddle/fluid/framework.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -373,8 +373,8 @@ def cuda_pinned_places(device_count=None):
373373
assert core.is_compiled_with_cuda(), \
374374
"Not compiled with CUDA"
375375
if device_count is None:
376-
device_count = _cpu_num()
377-
return [core.cuda_pinned_places()] * device_count
376+
device_count = len(_cuda_ids())
377+
return [core.CUDAPinnedPlace()] * device_count
378378

379379

380380
class NameScope(object):

python/paddle/fluid/tests/unittests/test_tensor.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -275,6 +275,10 @@ def test_tensor_poiter(self):
275275
self.assertTrue(
276276
isinstance(
277277
tensor._mutable_data(place, dtype), numbers.Integral))
278+
places = fluid.cuda_pinned_places()
279+
self.assertTrue(
280+
isinstance(
281+
tensor._mutable_data(places[0], dtype), numbers.Integral))
278282

279283

280284
if __name__ == '__main__':

0 commit comments

Comments
 (0)