File tree Expand file tree Collapse file tree 2 files changed +49
-0
lines changed Expand file tree Collapse file tree 2 files changed +49
-0
lines changed Original file line number Diff line number Diff line change 2222 'Event' ,
2323 'current_stream' ,
2424 'synchronize' ,
25+ 'device_count' ,
2526]
2627
2728
@@ -94,3 +95,25 @@ def synchronize(device=None):
9495 raise ValueError ("device type must be int or paddle.CUDAPlace" )
9596
9697 return core ._device_synchronize (device_id )
98+
99+
100+ def device_count ():
101+ '''
102+ Return the number of GPUs available.
103+
104+ Returns:
105+ int: the number of GPUs available.
106+
107+ Examples:
108+ .. code-block:: python
109+
110+ import paddle
111+
112+ paddle.device.cuda.device_count()
113+
114+ '''
115+
116+ num_gpus = core .get_cuda_device_count () if hasattr (
117+ core , 'get_cuda_device_count' ) else 0
118+
119+ return num_gpus
Original file line number Diff line number Diff line change 1+ # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
2+ #
3+ # Licensed under the Apache License, Version 2.0 (the "License");
4+ # you may not use this file except in compliance with the License.
5+ # You may obtain a copy of the License at
6+ #
7+ # http://www.apache.org/licenses/LICENSE-2.0
8+ #
9+ # Unless required by applicable law or agreed to in writing, software
10+ # distributed under the License is distributed on an "AS IS" BASIS,
11+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+ # See the License for the specific language governing permissions and
13+ # limitations under the License.
14+
15+ import paddle
16+ import unittest
17+
18+
19+ class TestDeviceCount (unittest .TestCase ):
20+ def test_device_count (self ):
21+ s = paddle .device .cuda .device_count ()
22+ self .assertIsNotNone (s )
23+
24+
25+ if __name__ == "__main__" :
26+ unittest .main ()
You can’t perform that action at this time.
0 commit comments