Update demoBERT C++ inference for triton

Signed-off-by: Rajeev Rao <rajeevrao@nvidia.com>
NVIDIA · May 11, 2021 · faced58 · faced58
1 parent 718c13d
commit faced58
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/demo/BERT/infer_c/bert_infer.h b/demo/BERT/infer_c/bert_infer.h
@@ -74,7 +74,7 @@ struct BertInference
             exit(-1);
         }
 
-        mEngine = TrtUniquePtr<ICudaEngine>(runtime->deserializeCudaEngine(bytes.data(), bytes.size(), nullptr));
+        mEngine = TrtUniquePtr<ICudaEngine>(runtime->deserializeCudaEngine(bytes.data(), bytes.size()));
         if (mEngine == nullptr)
         {
             gLogError << "Error deserializing CUDA engine\n";
@@ -175,7 +175,7 @@ struct BertInference
         {
             for (int i = 0; i < kBERT_INPUT_NUM; i++)
             {
-                mContext->setBindingDimensions(i + bindingIdxOffset, Dims2(mSeqLength, batchSize));
+                mContext->setBindingDimensions(i + bindingIdxOffset, Dims2(batchSize, mSeqLength));
             }
         }