diff --git a/paddle/fluid/operators/detection/yolo_box_op.cu b/paddle/fluid/operators/detection/yolo_box_op.cu index 01edf7b41b2a8dc6cce65a6e31d89025e884e58d..65dc73ef38323521590c9f5914ac13b321ef4469 100644 --- a/paddle/fluid/operators/detection/yolo_box_op.cu +++ b/paddle/fluid/operators/detection/yolo_box_op.cu @@ -15,7 +15,7 @@ limitations under the License. */ #include "paddle/fluid/memory/malloc.h" #include "paddle/fluid/operators/detection/yolo_box_op.h" #include "paddle/fluid/operators/math/math_function.h" - +#include "paddle/fluid/platform/gpu_launch_config.h" namespace paddle { namespace operators { @@ -108,11 +108,11 @@ class YoloBoxOpCUDAKernel : public framework::OpKernel { math::SetConstant set_zero; set_zero(dev_ctx, boxes, static_cast(0)); set_zero(dev_ctx, scores, static_cast(0)); + platform::GpuLaunchConfig config = + platform::GetGpuLaunchConfig1D(ctx.cuda_device_context(), n * box_num); - int grid_dim = (n * box_num + 512 - 1) / 512; - grid_dim = grid_dim > 8 ? 8 : grid_dim; - - KeYoloBoxFw<<>>( + KeYoloBoxFw<<>>( input_data, imgsize_data, boxes_data, scores_data, conf_thresh, anchors_data, n, h, w, an_num, class_num, box_num, input_size_h, input_size_w, clip_bbox, scale, bias);