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 9163720 commit 9e1a76fCopy full SHA for 9e1a76f
py/torch_tensorrt/dynamo/conversion/impl/embedding.py
@@ -94,7 +94,13 @@ 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)
+
98
+ # Notes: here offsets should always be 1d array
99
+ if len(offsets.shape) != 1:
100
+ raise TypeError(
101
+ f"The offsets should be in 1d array, here offset shape is {offsets.shape}."
102
+ )
103
+ offsets[-1] = len_embed
104
else:
105
# add the end index to offsets
106
offsets = np.append(offsets, len_embed)
0 commit comments