Browse Source

fix(mgb/version): fix nvinfer.h not found

GitOrigin-RevId: 981dd9a7d2
HuaHua404-patch-4
Megvii Engine Team 2 years ago
parent
commit
3fbceb3a38
1 changed files with 22 additions and 10 deletions
  1. +22
    -10
      src/core/impl/version.cpp

+ 22
- 10
src/core/impl/version.cpp View File

@@ -26,27 +26,39 @@ Version mgb::get_version() {
#endif
}

#if MGB_CUDA
#if __has_include("NvInfer.h")
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wnon-virtual-dtor"
#pragma GCC diagnostic ignored "-Wdeprecated-declarations"
#include "NvInfer.h"
int mgb::get_tensorrt_version() {
return NV_TENSORRT_VERSION;
}
#pragma GCC diagnostic pop
#else
int mgb::get_tensorrt_version() {
return -1;
}
#endif

#if __has_include("cuda.h")
#include "cuda.h"
#include "cudnn.h"
int mgb::get_cuda_version() {
return CUDA_VERSION;
}
int mgb::get_cudnn_version() {
return CUDNN_VERSION;
}
int mgb::get_tensorrt_version() {
return NV_TENSORRT_VERSION;
}
#else
int mgb::get_cuda_version() {
return -1;
}
#endif

#if __has_include("cudnn.h")
#include "cudnn.h"
int mgb::get_cudnn_version() {
return -1;
return CUDNN_VERSION;
}
int mgb::get_tensorrt_version() {
#else
int mgb::get_cudnn_version() {
return -1;
}
#endif


Loading…
Cancel
Save