From 1645b5efbfda753cb30db8596470834cf5fbb3c7 Mon Sep 17 00:00:00 2001 From: Evan Li Date: Wed, 15 Nov 2023 17:36:25 -0800 Subject: [PATCH] add bool type to convert_binary_elementwise --- py/torch_tensorrt/dynamo/conversion/impl/elementwise/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/elementwise/base.py b/py/torch_tensorrt/dynamo/conversion/impl/elementwise/base.py index 3700242fe7..7fe5c99eb2 100644 --- 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