Merge pull request #902 from NVIDIA/one_hot

peri044 · web-flow · commit 8801573d1800 · 2022-04-05T00:53:41.000-07:00
fix(//core): Take user setting in the case we can't determine the
diff --git a/core/compiler.cpp b/core/compiler.cpp
@@ -344,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

Original file line number	Diff line number	Diff line change
`@@ -344,9 +344,10 @@ void MapInputsAndDetermineDTypes(`
`344`	`344`	`ss << "- Disable partial compilation by setting require_full_compilation to True";`
`345`	`345`	`auto warn_str = ss.str();`
`346`	`346`	`LOG_WARNING(warn_str);`
`347`		`- // Overwrite type map with user settings`
`348`		`- first_use_type_map[in] = {util::TRTDataTypeToScalarType(cfg.convert_info.inputs.find(in)->second.dtype)};`
`349`	`347`	`}`
	`348`	`+ // Overwrite type map with user settings`
	`349`	`+ // We use this map for partitiioning since we need c10::ScalarTypes not nvinfer::DataTypes`
	`350`	`+ first_use_type_map[in] = {util::TRTDataTypeToScalarType(cfg.convert_info.inputs.find(in)->second.dtype)};`
`350`	`351`	`}`
`351`	`352`	`} else {`
`352`	`353`	`// The user defined the type so no changes are necessary`