blob: 544b48a679807c0a4858c01d3205de5b3387fca5 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
|
#include "tensorflow/core/util/use_cudnn.h"
#include <stdlib.h>
#include "tensorflow/core/platform/port.h"
namespace tensorflow {
bool CanUseCudnn() {
const char* tf_use_cudnn = getenv("TF_USE_CUDNN");
if (tf_use_cudnn != nullptr) {
string tf_use_cudnn_str = tf_use_cudnn;
if (tf_use_cudnn_str == "0") {
return false;
}
}
return true;
}
} // namespace tensorflow
|
