mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

Merge branch 'ceccocats:master' into master

Browse Source
This commit is contained in:
Adam Cuellar
2021-11-09 16:50:40 -05:00
committed by GitHub
parent 46daa78ca9 d6fb6c6af4
commit 31af5396bb
1 changed files with 9 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/kernels/postprocessing.cu
+9 -4
View File
@@ -46,11 +46,16 @@ void maxElem_kernel(float *src_begin, float *dst_begin, const int n_classes, con
if (i > size)
return;
thrust::device_ptr<float> dPbeg ( &src_begin[i*n_classes] ) ;
thrust::device_ptr<float> dPend = dPbeg + n_classes;
thrust::device_ptr<float> result = thrust::max_element(thrust::device,dPbeg, dPend);
float max = 0;
int max_idx = 0;
for( int j = i*n_classes; j < i*n_classes + n_classes; ++j ){
if( src_begin[j] > max ){
max = src_begin[j];
max_idx = j;
}
}
dst_begin[i] = result - dPbeg;
dst_begin[i] = max_idx - i*n_classes;
}
void maxElem(dnnType *src_begin, dnnType *dst_begin, const int c, const int h, const int w){