Skip to content

Commit bc2f56f

Browse files
committed
Updating the type for the inputs to the select layer
1 parent 88e5601 commit bc2f56f

File tree

1 file changed

+3
-3
lines changed
  • py/torch_tensorrt/dynamo/conversion/impl/condition

1 file changed

+3
-3
lines changed

py/torch_tensorrt/dynamo/conversion/impl/condition/ops.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -104,9 +104,9 @@ def select(
104104
target: Target,
105105
source_ir: Optional[SourceIR],
106106
name: str,
107-
input: Union[TRTTensor],
108-
other: Union[TRTTensor],
109-
condition: Union[TRTTensor],
107+
input: TRTTensor,
108+
other: TRTTensor,
109+
condition: TRTTensor,
110110
) -> TRTTensor:
111111
select_layer = ctx.net.add_select(condition, input, other)
112112
set_layer_name(select_layer, target, name + "_select", source_ir)

0 commit comments

Comments
 (0)