fix: Fix how ITensorList is detected

peri044 · peri044 · commit 76dc8046972c · 2023-04-05T15:34:18.000-07:00
Signed-off-by: Dheeraj Peri &lt;peri.dheeraj@gmail.com&gt;
diff --git a/core/conversion/var/Var.cpp b/core/conversion/var/Var.cpp
@@ -147,11 +147,15 @@ bool Var::isITensor() const {
 }
 
 bool Var::isITensorList() {
-  if (isList() && ptr_.ivalue->isCustomClass()) {
-    return true;
-  } else {
-    return false;
+  // Unpack the Var as a List and check if each entry is a custom class since
+  // ITensors are stored in CustomClassHolder
+  auto ival_list = ptr_.ivalue->toList();
+  for (int i = 0; i < ival_list.size(); i++) {
+    if (!ival_list.get(i).isCustomClass()) {
+      return false;
+    }
   }
+  return true;
 }
 
 std::vector<nvinfer1::ITensor*> Var::unwrapToITensorList() {

Original file line number	Diff line number	Diff line change
`@@ -147,11 +147,15 @@ bool Var::isITensor() const {`
`147`	`147`	`}`
`148`	`148`
`149`	`149`	`bool Var::isITensorList() {`
`150`		`- if (isList() && ptr_.ivalue->isCustomClass()) {`
`151`		`- return true;`
`152`		`- } else {`
`153`		`- return false;`
	`150`	`+ // Unpack the Var as a List and check if each entry is a custom class since`
	`151`	`+ // ITensors are stored in CustomClassHolder`
	`152`	`+ auto ival_list = ptr_.ivalue->toList();`
	`153`	`+ for (int i = 0; i < ival_list.size(); i++) {`
	`154`	`+ if (!ival_list.get(i).isCustomClass()) {`
	`155`	`+ return false;`
	`156`	`+ }`
`154`	`157`	`}`
	`158`	`+ return true;`
`155`	`159`	`}`
`156`	`160`
`157`	`161`	`std::vector<nvinfer1::ITensor*> Var::unwrapToITensorList() {`