File tree Expand file tree Collapse file tree 1 file changed +9
-3
lines changed
py/torch_tensorrt/dynamo/conversion/impl Expand file tree Collapse file tree 1 file changed +9
-3
lines changed Original file line number Diff line number Diff line change @@ -46,9 +46,15 @@ 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 ((x .numel (),), np .array ([x ], dtype = np .int32 ))
50- set_layer_name (t , target , f"{ name } _dim{ i } _const" )
51- t = t .get_output (0 )
49+ const_arr = np .array ([x ], dtype = np .int32 )
50+ shape = (1 ,)
51+ if not isinstance (x , int ):
52+ const_arr = np .array (x , dtype = np .int32 )
53+ shape = (x .numel (),)
54+
55+ layer = ctx .net .add_constant (shape , const_arr )
56+ set_layer_name (layer , target , f"{ name } _dim{ i } _const" )
57+ t = layer .get_output (0 )
5258
5359 # optional cast
5460 if cast_dtype and isinstance (t , TRTTensor ):
You can’t perform that action at this time.
0 commit comments