Fix scatternd bug

Longzhong-Lin · May 1, 2021 · cebbe9a · cebbe9a
1 parent aeba71e
commit cebbe9a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tensorrt/samples/centerpoint/ScatterNDPlugin.cu b/tensorrt/samples/centerpoint/ScatterNDPlugin.cu
@@ -120,7 +120,7 @@ int ScatterNDPlugin::enqueue(int batchSize, const void* const* inputs, void** ou
     int totalElems = mOutputSize[0]*channel_num;
 
     dim3 blockSize(THREAD_NUM);
-    dim3 gridsize(max_index_num/channel_num/blockSize.x);
+    dim3 gridsize(max_index_num/blockSize.x+1);
 
     switch (mDataType)
     {
@@ -311,4 +311,4 @@ IPluginV2Ext* ScatterNDSamplePluginCreator::deserializePlugin(const char* name,
     return new ScatterNDPlugin(name, serialData, serialLength);
 }
 
-REGISTER_TENSORRT_PLUGIN(ScatterNDSamplePluginCreator);
+REGISTER_TENSORRT_PLUGIN(ScatterNDSamplePluginCreator);