Merge pull request #255 from NVIDIA/fix_dtype

Support bool datatype of input/output tensors

Author: narendasan

core/util/trt_util.cpp

@@ -229,6 +229,7 @@ const std::unordered_map<at::ScalarType, nvinfer1::DataType>& get_at_trt_type_ma
       {at::kHalf, nvinfer1::DataType::kHALF},
       {at::kInt, nvinfer1::DataType::kINT32},
       {at::kChar, nvinfer1::DataType::kINT8},
+      {at::kBool, nvinfer1::DataType::kBOOL},
   };
   return at_trt_type_map;
 }
@@ -239,6 +240,7 @@ const std::unordered_map<nvinfer1::DataType, at::ScalarType>& get_trt_at_type_ma
       {nvinfer1::DataType::kHALF, at::kHalf},
       {nvinfer1::DataType::kINT32, at::kInt},
       {nvinfer1::DataType::kINT8, at::kChar},
+      {nvinfer1::DataType::kBOOL, at::kBool},
   };
   return trt_at_type_map;
 }
@@ -249,15 +251,19 @@ const std::unordered_map<nvinfer1::DataType, at::ScalarType>& get_trt_aten_type_
 }
 
 at::ScalarType toATenDType(nvinfer1::DataType t) {
-  return get_trt_aten_type_map().at(t);
+  auto trt_aten_type_map = get_trt_aten_type_map();
+  TRTORCH_CHECK(trt_aten_type_map.find(t) != trt_aten_type_map.end(), "Unsupported TensorRT datatype");
+  return trt_aten_type_map.at(t);
 }
 
 const std::unordered_map<at::ScalarType, nvinfer1::DataType>& get_aten_trt_type_map() {
   return get_at_trt_type_map();
 }
 
 nvinfer1::DataType toTRTDataType(at::ScalarType t) {
-  return get_aten_trt_type_map().at(t);
+  auto aten_trt_type_map = get_aten_trt_type_map();
+  TRTORCH_CHECK(aten_trt_type_map.find(t) != aten_trt_type_map.end(), "Unsupported Aten datatype");
+  return aten_trt_type_map.at(t);
 }
 
 c10::optional<nvinfer1::DataType> toTRTDataType(caffe2::TypeMeta dtype) {

tests/util/util.cpp

@@ -17,7 +17,9 @@ bool checkRtol(const at::Tensor& diff, const std::vector<at::Tensor> inputs, flo
 
 bool almostEqual(const at::Tensor& a, const at::Tensor& b, float threshold) {
   LOG_DEBUG(a << std::endl << b << std::endl);
-  return checkRtol(a - b, {a, b}, threshold);
+  auto a_float = a.toType(at::kFloat);
+  auto b_float = b.toType(at::kFloat);
+  return checkRtol(a_float - b_float, {a_float, b_float}, threshold);
 }
 
 bool exactlyEqual(const at::Tensor& a, const at::Tensor& b) {