add cudnn data type processing for ATen tensor (pytorch#2087)

Author: killeent

torch/csrc/cudnn/Types.cpp

@@ -34,6 +34,19 @@ cudnnDataType_t getCudnnDataType(const thpp::Tensor& tensor)
   throw std::runtime_error(msg);
 }
 
+cudnnDataType_t getCudnnDataType(const at::Tensor& tensor) {
+  if (tensor.type().scalarType() == at::kFloat) {
+    return CUDNN_DATA_FLOAT;
+  } else if (tensor.type().scalarType() == at::kDouble) {
+    return CUDNN_DATA_DOUBLE;
+  } else if (tensor.type().scalarType() == at::kHalf) {
+    return CUDNN_DATA_HALF;
+  }
+  std::string msg("getCudnnDataType() not supported for ");
+  msg += at::toString(tensor.type().scalarType());
+  throw std::runtime_error(msg);
+}
+
 PyObject * getTensorClass(PyObject *args)
 {
   for (int i = 0; i < PyTuple_Size(args); i++) {

torch/csrc/cudnn/Types.h

@@ -7,12 +7,14 @@
 #include <cudnn.h>
 #include "../Types.h"
 #include <THPP/THPP.h>
+#include <ATen/Tensor.h>
 
 namespace torch { namespace cudnn {
 
 PyObject * getTensorClass(PyObject *args);
 cudnnDataType_t getCudnnDataType(PyObject *tensorClass);
 cudnnDataType_t getCudnnDataType(const thpp::Tensor& tensor);
+cudnnDataType_t getCudnnDataType(const at::Tensor& tensor);
 void _THVoidTensor_assertContiguous(THVoidTensor *tensor, const std::string& name);
 
 #define THVoidTensor_assertContiguous(tensor) \