From f632706c18ee926700ad3fbf73d4952ed648c395 Mon Sep 17 00:00:00 2001
From: Dong Zhihong <dzhwinter@gmail.com>
Date: Thu, 26 Oct 2017 15:09:14 -0700
Subject: [PATCH] fix based on comment

---
 paddle/pybind/pybind.cc                               | 2 ++
 python/paddle/v2/framework/tests/test_nccl_init_op.py | 7 +++----
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/paddle/pybind/pybind.cc b/paddle/pybind/pybind.cc
index 9288468a038..35fbf4d04a6 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 8aed14c15d9..03d46d1c600 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}
-- 
GitLab