refactor(//core): Unify the dtype work for the partition section

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

core/compiler.cpp

@@ -329,7 +329,6 @@ void MapInputsAndDetermineDTypes(
       } 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");
-          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;
@@ -345,9 +344,10 @@ void MapInputsAndDetermineDTypes(
             ss << "- Disable partial compilation by setting require_full_compilation to True";
             auto warn_str = ss.str();
             LOG_WARNING(warn_str);
-            // Overwrite type map with user settings
-            first_use_type_map[in] = {util::TRTDataTypeToScalarType(cfg.convert_info.inputs.find(in)->second.dtype)};
           }
+          // Overwrite type map with user settings
+          // We use this map for partitiioning since we need c10::ScalarTypes not nvinfer::DataTypes
+          first_use_type_map[in] = {util::TRTDataTypeToScalarType(cfg.convert_info.inputs.find(in)->second.dtype)};
         }
       } else {
         // The user defined the type so no changes are necessary