validator for index non zero case- else error- Could not find any implementation for node [NON_ZERO]-[aten_ops.index.Tensor]-[index_tensor_bool_nonzero_1]

Author: apbose

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -9,7 +9,7 @@
 from tensorrt import ITensor as TRTTensor
 from torch.fx.node import Argument, Node, Target
 from torch_tensorrt._features import needs_not_tensorrt_rtx
-from torch_tensorrt._utils import is_tensorrt_version_supported
+from torch_tensorrt._utils import is_tensorrt_version_supported, is_thor
 from torch_tensorrt.dynamo._settings import CompilationSettings
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.dynamo.conversion import impl
@@ -24,7 +24,8 @@
     get_positive_dim,
     is_only_operator_on_placeholder,
 )
-from torch_tensorrt.dynamo.utils import is_thor
+
+from torch_tensorrt._utils import is_thor
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -425,9 +426,24 @@ def index_dtype_validator(
     return True
 
 
+def index_nonbool_validator(
+    node: Node, settings: Optional[CompilationSettings] = None
+) -> bool:
+    # for thor, we don't support boolean indices
+    if is_thor():
+        index = node.args[1]
+        for ind in index:
+            if ind is not None:
+                val = ind.meta.get("val")
+                if val is not None and val.dtype == torch.bool:
+                    return False
+    return True
+
+
 @dynamo_tensorrt_converter(
     torch.ops.aten.index.Tensor,
-    capability_validator=index_dtype_validator,
+    capability_validator=lambda node, settings: index_dtype_validator(node, settings)
+    and index_nonbool_validator(node, settings),
     supports_dynamic_shapes=True,
     requires_output_allocator=True,
 )