Fix cuda sum reduce
This commit is contained in:
parent
a6da5bc9af
commit
9981bb8b6f
8 changed files with 12 additions and 10 deletions
|
|
@ -171,9 +171,9 @@ extern "C" {
|
|||
if (strcmp(tensor->device, "cuda") == 0) {
|
||||
|
||||
float* result_data;
|
||||
int num_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
|
||||
cudaMalloc((void **)&num_of_blocks, 1 * sizeof(float));
|
||||
sum_tensor_cuda(tensor, result_data, int num_of_blocks);
|
||||
int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
|
||||
cudaMalloc((void**)&result_data, number_of_blocks * sizeof(float));
|
||||
sum_tensor_cuda(tensor, result_data, number_of_blocks);
|
||||
return create_tensor(result_data, shape, ndim, device);
|
||||
}
|
||||
else {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue