Fix a bug that causes error when batchSize > 1 (#441)

Signed-off-by: Jeffery Zhao <zeyuz@nvidia.com>
NVIDIA · Mar 19, 2020 · 9d37f5d · 9d37f5d
1 parent f5e8b8c
commit 9d37f5d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/plugin/batchTilePlugin/batchTilePlugin.cpp b/plugin/batchTilePlugin/batchTilePlugin.cpp
@@ -95,7 +95,7 @@ DataType BatchTilePlugin::getOutputDataType(int index, const nvinfer1::DataType*
 
 int BatchTilePlugin::enqueue(int batchSize, const void* const* inputs, void** outputs, void*, cudaStream_t stream)
 {
-    float* output = reinterpret_cast<float*>(outputs[0]);
+    char* output = reinterpret_cast<char*>(outputs[0]);
     // expand to batch size
     for (int i = 0; i < batchSize; i++)
     {
@@ -210,4 +210,4 @@ IPluginV2Ext* BatchTilePluginCreator::createPlugin(const char* name, const Plugi
 IPluginV2Ext* BatchTilePluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
 {
     return new BatchTilePlugin(name, serialData, serialLength);
-}
+}