|
|
@@ -267,10 +267,9 @@ static __global__ void softmax(float *predict, int length)
|
|
|
}
|
|
|
|
|
|
static void classfier_softmax(float *predict,int length, cudaStream_t stream) {
|
|
|
- dim3 grid((dst_width + 31) / 32, (dst_height + 31) / 32);
|
|
|
- dim3 block(32, 32);
|
|
|
+ int block_size = 256;
|
|
|
|
|
|
- checkKernel(softmax<<<grid, block, 0, stream>>>(predict, length));
|
|
|
+ checkKernel(softmax<<<1, block_size, block_size * sizeof(float), stream>>>(predict, length));
|
|
|
}
|
|
|
|
|
|
class InferImpl : public Infer {
|
