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

fix

Browse Source
This commit is contained in:
Francesco Gatti
2017-08-04 16:16:03 +00:00
parent b20a2e2902
commit 34198a4e8d
6 changed files with 55 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/NetworkRT.cpp
+10 -7
View File
@@ -10,12 +10,10 @@ using namespace nvinfer1;
#include "pluginsRT/RegionRT.cpp"
// Logger for info/warning/errors
class Logger : public ILogger
{
void log(Severity severity, const char* msg) override
{
std::cout <<"TENSORRT: "<< msg << std::endl;
}
class Logger : public ILogger {
void log(Severity severity, const char* msg) override {
std::cout <<"TENSORRT LOG: "<< msg << std::endl;
}
} loggerRT;
namespace tkDNN {
@@ -24,6 +22,11 @@ std::map<Layer*, nvinfer1::ITensor*>tensors;
NetworkRT::NetworkRT(Network *net) {
float rt_ver = float(NV_TENSORRT_MAJOR) +
float(NV_TENSORRT_MINOR)/10 +
float(NV_TENSORRT_PATCH)/100;
std::cout<<"New NetworkRT (TensorRT v"<<rt_ver<<")\n";
builderRT = createInferBuilder(loggerRT);
networkRT = builderRT->createNetwork();
dtRT = DataType::kFLOAT;
@@ -250,4 +253,4 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Region *l) {
return lRT->getOutput(0);
}
}
}