fix limit_by_capacity op

PaddlePaddle · Sep 27, 2021 · 0f13c5d · 0f13c5d · paddle-bot-old · Sep 27, 2021
1 parent da99d23
commit 0f13c5d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 9 deletions.
diff --git a/paddle/fluid/operators/limit_by_capacity_op.cc b/paddle/fluid/operators/limit_by_capacity_op.cc
@@ -52,11 +52,8 @@ class LimitByCapacityOpMaker : public framework::OpProtoAndCheckerMaker {
 namespace ops = paddle::operators;
 namespace plat = paddle::platform;
 
-REGISTER_OP_WITHOUT_GRADIENT(LimitByCapacityOp, ops::LimitByCapacityOp,
-                             ops::LimitByCapacityOpMaker);
-
-REGISTER_OPERATOR(limit_by_capacity, ops::LimitByCapacityOp,
-                  ops::LimitByCapacityOpMaker);
-
 REGISTER_OP_CPU_KERNEL(limit_by_capacity, ops::LimitByCapacityOpCPUKernel<int>,
                        ops::LimitByCapacityOpCPUKernel<int64_t>);
+
+REGISTER_OP_WITHOUT_GRADIENT(limit_by_capacity, ops::LimitByCapacityOp,
+                             ops::LimitByCapacityOpMaker);
diff --git a/paddle/fluid/operators/limit_by_capacity_op.cu b/paddle/fluid/operators/limit_by_capacity_op.cu
@@ -26,8 +26,8 @@ using LoDTensor = framework::LoDTensor;
 using Tensor = framework::Tensor;
 
 template <typename T>
-__global__ void LimitByCapacity(const T* expc, int* cap, T* out,
-                                const int n_expert, const int n_worker) {
+__global__ void limit_by_capacity_impl(const T* expc, int* cap, T* out,
+                                       const int n_expert, const int n_worker) {
   int eid = blockIdx.y;
   int wid = blockIdx.x * blockDim.x + threadIdx.x;
   if (wid < n_worker) {
@@ -67,7 +67,7 @@ class LimitByCapacityOpCUDAKernel : public framework::OpKernel<T> {
     framework::TensorCopy(*capacity, place, dev_ctx, &capacity_copy);
     int* cap_data = capacity_copy.mutable_data<int>(place);
 
-    LimitByCapacity<T><<<grid_dim, block_dim, 0, dev_ctx.stream()>>>(
+    limit_by_capacity_impl<T><<<grid_dim, block_dim, 0, dev_ctx.stream()>>>(
         ec_data, cap_data, out_data, n_expert, n_worker);
   }
 };