Caffe2 - C++ API
A deep learning, cross platform ML framework
Types.cpp
1 #include <ATen/cudnn/Types.h>
2 
3 #include <ATen/ATen.h>
4 
5 namespace at { namespace native {
6 
7 cudnnDataType_t getCudnnDataType(const at::Tensor& tensor) {
8  if (tensor.scalar_type() == at::kFloat) {
9  return CUDNN_DATA_FLOAT;
10  } else if (tensor.scalar_type() == at::kDouble) {
11  return CUDNN_DATA_DOUBLE;
12  } else if (tensor.scalar_type() == at::kHalf) {
13  return CUDNN_DATA_HALF;
14  }
15  std::string msg("getCudnnDataType() not supported for ");
16  msg += toString(tensor.scalar_type());
17  throw std::runtime_error(msg);
18 }
19 
20 int64_t cudnn_version() {
21  return CUDNN_VERSION;
22 }
23 
24 }} // namespace at::cudnn
Flush-To-Zero and Denormals-Are-Zero mode.