mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
27 lines
707 B
C++
27 lines
707 B
C++
#include "caffe2/core/common_cudnn.h"
|
|
#include "caffe2/core/cudnn_wrappers.h"
|
|
|
|
#include "caffe2/core/init.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
CuDNNWrapper::PerGPUCuDNNStates& CuDNNWrapper::cudnn_states() {
|
|
// New it (never delete) to avoid calling the destructors on process
|
|
// exit and racing against the CUDA shutdown sequence.
|
|
static auto* p = new CuDNNWrapper::PerGPUCuDNNStates();
|
|
CHECK_NOTNULL(p);
|
|
return *p;
|
|
}
|
|
|
|
namespace {
|
|
bool PrintCuDNNInfo(int*, char***) {
|
|
VLOG(1) << "Caffe2 is built with CuDNN version " << CUDNN_VERSION;
|
|
return true;
|
|
}
|
|
|
|
REGISTER_CAFFE2_INIT_FUNCTION(PrintCuDNNInfo, &PrintCuDNNInfo,
|
|
"Print CuDNN Info.");
|
|
|
|
} // namespace
|
|
} // namespace caffe2
|