File size: 594 Bytes
4616d53
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
#include <NvInfer.h>
#include <cstdint>
#include <ostream>
// Off-path stubs not provided by the util files we link.
namespace onnx2trt {
std::ostream& operator<<(std::ostream& o, nvinfer1::Dims64 const&){ return o << "<dims>"; }
bool convertDtype(int32_t onnx_dtype, nvinfer1::DataType* trt){
    switch(onnx_dtype){
        case 1: case 11: *trt=nvinfer1::DataType::kFLOAT; return true;
        case 10: *trt=nvinfer1::DataType::kHALF; return true;
        case 6: case 7: *trt=nvinfer1::DataType::kINT32; return true;
        default: *trt=nvinfer1::DataType::kFLOAT; return true;
    }
}
}