File tree Expand file tree Collapse file tree 2 files changed +6
-2
lines changed Expand file tree Collapse file tree 2 files changed +6
-2
lines changed Original file line number Diff line number Diff line change @@ -373,8 +373,8 @@ def cuda_pinned_places(device_count=None):
373
373
assert core .is_compiled_with_cuda (), \
374
374
"Not compiled with CUDA"
375
375
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
378
378
379
379
380
380
class NameScope (object ):
Original file line number Diff line number Diff line change @@ -275,6 +275,10 @@ def test_tensor_poiter(self):
275
275
self .assertTrue (
276
276
isinstance (
277
277
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 ))
278
282
279
283
280
284
if __name__ == '__main__' :
You can’t perform that action at this time.
0 commit comments