File tree Expand file tree Collapse file tree 1 file changed +2
-2
lines changed
py/torch_tensorrt/dynamo/conversion/impl Expand file tree Collapse file tree 1 file changed +2
-2
lines changed Original file line number Diff line number Diff line change @@ -184,7 +184,7 @@ def embedding_bag_with_ITensor_offsets(
184
184
loop1 = ctx .net .add_loop ()
185
185
trip_limit1 = ctx .net .add_constant (
186
186
shape = (),
187
- weights = trt .Weights (np .array ([offsets .shape [0 ] - 1 ], dtype = np .dtype ( "i" ) )),
187
+ weights = trt .Weights (np .array ([offsets .shape [0 ] - 1 ], dtype = np .int32 )),
188
188
).get_output (0 )
189
189
loop1 .add_trip_limit (trip_limit1 , trt .TripLimit .COUNT )
190
190
@@ -205,7 +205,7 @@ def embedding_bag_with_ITensor_offsets(
205
205
###### Inner loop: traverse indices ######
206
206
loop2 = ctx .net .add_loop ()
207
207
trip_limit2 = ctx .net .add_constant (
208
- shape = (), weights = trt .Weights (np .array ([len_embed ], dtype = np .dtype ( "i" ) ))
208
+ shape = (), weights = trt .Weights (np .array ([len_embed ], dtype = np .int32 ))
209
209
).get_output (0 )
210
210
loop2 .add_trip_limit (trip_limit2 , trt .TripLimit .COUNT )
211
211
rec2_j_tensor = loop2 .add_recurrence (constant_0 )
You can’t perform that action at this time.
0 commit comments