fix: Use user provided dtype when we can't infer it from the graph

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

Diff for: core/compiler.cpp

@@ -328,7 +328,8 @@ void MapInputsAndDetermineDTypes(
         spec.dtype = nvinfer1::DataType::kFLOAT;
       } else if (spec.dtype_is_user_defined && cfg.partition_info.enabled) {
         if (!est_type_opt) {
-          LOG_INFO("Cannot infer input tensor dtype in graph, unable to verify user input dtype settings");
+          LOG_INFO("Cannot infer input tensor dtype in graph. Using user provided input dtype settings");
+          first_use_type_map[in] = {util::TRTDataTypeToScalarType(cfg.convert_info.inputs.find(in)->second.dtype)};
         } else {
           if (util::TRTDataTypeToScalarType(cfg.convert_info.inputs.find(in)->second.dtype) != est_type_opt.value()) {
             std::stringstream ss;

Diff for: core/ir/Input.cpp

@@ -40,6 +40,13 @@ bool valid_dtype_format_combo(nvinfer1::DataType dtype, nvinfer1::TensorFormat f
         default:
           return false;
       }
+    case nvinfer1::DataType::kBOOL: // Supports Linear (NCHW)
+      switch (format) {
+        case nvinfer1::TensorFormat::kLINEAR:
+          return true;
+        default:
+          return false;
+      }
     default:
       return false;
   }
@@ -48,7 +55,7 @@ bool valid_dtype_format_combo(nvinfer1::DataType dtype, nvinfer1::TensorFormat f
 bool valid_input_dtype(nvinfer1::DataType dtype) {
   switch (dtype) {
     case nvinfer1::DataType::kBOOL:
-      return false;
+      return true;
     case nvinfer1::DataType::kFLOAT:
       return true;
     case nvinfer1::DataType::kHALF:
@@ -153,4 +160,4 @@ std::ostream& operator<<(std::ostream& os, const Input& input) {
 
 } // namespace ir
 } // namespace core
-} // namespace torch_tensorrt
+} // namespace torch_tensorrt

Diff for: py/setup.py

@@ -239,8 +239,6 @@ def run(self):
         libraries=["torchtrt"],
         include_dirs=[
             dir_path + "torch_tensorrt/csrc", dir_path + "torch_tensorrt/include",
-            dir_path + "/../bazel-TRTorch/external/tensorrt/include",
-            dir_path + "/../bazel-Torch-TensorRT-Preview/external/tensorrt/include",
             dir_path + "/../bazel-Torch-TensorRT/external/tensorrt/include", dir_path + "/../"
         ],
         extra_compile_args=[

Diff for: py/torch_tensorrt/csrc/torch_tensorrt_py.cpp

@@ -185,7 +185,7 @@ PYBIND11_MODULE(_C, m) {
       .value("float16", DataType::kHalf, "16 bit floating point number")
       .value("int8", DataType::kChar, "8 bit integer number")
       .value("int32", DataType::kInt32, "32 bit integer number")
-      .value("bool", DataType::kChar, "Boolean value")
+      .value("bool", DataType::kBool, "Boolean value")
       .value("unknown", DataType::kUnknown, "Unknown data type")
       .export_values();