Skip to content

Commit 8db4c74

Browse files
committed
adding torch Tensor in check condition for CI
1 parent 975b403 commit 8db4c74

File tree

1 file changed

+1
-1
lines changed
  • py/torch_tensorrt/dynamo/conversion/impl

1 file changed

+1
-1
lines changed

py/torch_tensorrt/dynamo/conversion/impl/shape.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -143,7 +143,7 @@ def to_trt_shape_tensor(
143143
trt_tensors = []
144144

145145
for i, s in enumerate(shape_list):
146-
if isinstance(s, int):
146+
if isinstance(s, (int, torch.Tensor)):
147147
const = ctx.net.add_constant((1,), np.array([s], dtype=np.int32))
148148
set_layer_name(const, target, f"{name}_dim{i}_const")
149149
trt_tensors.append(const.get_output(0))

0 commit comments

Comments
 (0)