Update yoloPlugins.cpp

This commit is contained in:
Marcos Luciano
2022-02-18 03:51:08 -03:00
committed by GitHub
parent 329daeb150
commit 0c815aef80

View File

@@ -212,8 +212,8 @@ int YoloLayer::enqueue(
} }
else { else {
void* softmax; void* softmax;
cudaMallocHost(&softmax, sizeof(outputs[0])); CHECK(cudaMallocHost(&softmax, sizeof(outputs[0])));
cudaMemcpy(softmax, outputs[0], sizeof(outputs[0]), cudaMemcpyHostToDevice); CHECK(cudaMemcpy(softmax, outputs[0], sizeof(outputs[0]), cudaMemcpyHostToDevice));
CHECK(cudaYoloLayer_v2( CHECK(cudaYoloLayer_v2(
inputs[0], outputs[0], softmax, batchSize, m_GridSizeX, m_GridSizeY, m_NumClasses, m_NumBoxes, inputs[0], outputs[0], softmax, batchSize, m_GridSizeX, m_GridSizeY, m_NumClasses, m_NumBoxes,