We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent 9643560 commit e83725eCopy full SHA for e83725e
py/torch_tensorrt/dynamo/conversion/impl/embedding.py
@@ -94,7 +94,12 @@ def embedding_bag_with_traversable_offsets(
94
# however, pytorch doc says if `include_last_offset` is True, the size of offsets
95
# is equal to the number of bags + 1. The last element is the size of the input,
96
# or the ending index position of the last bag (sequence).
97
- offsets.itemset(-1, len_embed)
+ # Notes: here offsets should always be 1d array
98
+ if len(offsets.shape) != 1:
99
+ raise TypeError(
100
+ f"The offsets should be in 1d array, here offset shape is {offsets.shape}."
101
+ )
102
+ offsets[-1] = len_embed
103
else:
104
# add the end index to offsets
105
offsets = np.append(offsets, len_embed)
0 commit comments