add bool type to convert_binary_elementwise

pytorch · Nov 16, 2023 · 1645b5e · 1645b5e
1 parent af49b8b
commit 1645b5e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/elementwise/base.py b/py/torch_tensorrt/dynamo/conversion/impl/elementwise/base.py
@@ -58,8 +58,8 @@ def convert_binary_elementwise(
     source_ir: Optional[SourceIR],
     name: str,
     op_type: trt.ElementWiseOperation,
-    lhs_val: Union[int, float, TRTTensor, torch.Tensor],
-    rhs_val: Union[int, float, TRTTensor, torch.Tensor],
+    lhs_val: Union[int, float, bool, TRTTensor, torch.Tensor],
+    rhs_val: Union[int, float, bool, TRTTensor, torch.Tensor],
 ) -> TRTTensor:
     """
     This function adds a TensorRT elementwise layer. We allow both operands to be