Skip to content

Commit ce56ca6

Browse files
committed
addressing embedding bag CI error
1 parent 5b77d3f commit ce56ca6

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/cat.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ def unify_and_concat_trt_tensors(
4646
elif isinstance(x, int) and not has_dynamic and not force_trt_output:
4747
t = x # pure static path
4848
else:
49-
t = ctx.net.add_constant((1,), np.array([x], dtype=np.int32))
49+
t = ctx.net.add_constant((x.numel(),), np.array([x], dtype=np.int32))
5050
set_layer_name(t, target, f"{name}_dim{i}_const")
5151
t = t.get_output(0)
5252

0 commit comments

Comments
 (0)