diff --git a/paddle/pybind/pybind.cc b/paddle/pybind/pybind.cc index 9288468a0381c5dbcf54255bef6e97aa53f510c4..35fbf4d04a6a199d01ee6c52e901d1721c33c482 100644 --- a/paddle/pybind/pybind.cc +++ b/paddle/pybind/pybind.cc @@ -34,6 +34,7 @@ limitations under the License. */ #ifdef PADDLE_WITH_CUDA #include "paddle/operators/nccl/nccl_gpu_common.h" +#include "paddle/platform/gpu_info.h" #endif namespace paddle { @@ -482,6 +483,7 @@ All parameter, weight, gradient are variables in Paddle. BindOpDesc(m); m.def("op_support_gpu", OpSupportGPU); + m.def("get_cuda_device_count", platform::GetCUDADeviceCount); return m.ptr(); } diff --git a/python/paddle/v2/framework/tests/test_nccl_init_op.py b/python/paddle/v2/framework/tests/test_nccl_init_op.py index 8aed14c15d9e99e7a964117f64a8cbe532b57ac8..03d46d1c600afc40de0ee5d28d8259135999436a 100644 --- a/python/paddle/v2/framework/tests/test_nccl_init_op.py +++ b/python/paddle/v2/framework/tests/test_nccl_init_op.py @@ -5,11 +5,10 @@ from paddle.v2.framework.op import Operator import paddle.v2.framework.core as core from op_test import OpTest, create_op, set_input -gpu_list = "0,1,2,3" - -if not core.is_compile_gpu() or not gpu_list: +if not core.is_compile_gpu(): exit(0) +gpu_count = core.get_cuda_device_count g_scope = core.Scope() g_ctx = core.DeviceContext.create(core.CPUPlace()) @@ -17,7 +16,7 @@ g_ctx = core.DeviceContext.create(core.CPUPlace()) class TestNCCLInit(unittest.TestCase): def test_init(self): self.op_type = "ncclInit" - self.gpus = [int(g) for g in gpu_list.split(",")] + self.gpus = [int(g) for g in range(gpu_count)] self.inputs = {} self.attrs = {"gpus": self.gpus}