diff --git a/py/torch_tensorrt/dynamo/conversion/impl/embedding.py b/py/torch_tensorrt/dynamo/conversion/impl/embedding.py index 0a723618eb..b856023f8e 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/embedding.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/embedding.py @@ -30,12 +30,13 @@ def embedding( ) -> TRTTensor: indices_tensor = input embedding_tensor = weight - if isinstance(indices_tensor, torch.Tensor) and indices_tensor.dtype == torch.int64: - raise RuntimeError( - "The `embedding` op has indices_tensor dtype=int64. This is incorrect since it has to be int32 to run on TRT." - ) indices_tensor = get_trt_tensor(ctx, indices_tensor, f"{name}_indices_tensor") embedding_tensor = get_trt_tensor(ctx, embedding_tensor, f"{name}_embedding_tensor") + if indices_tensor.dtype not in (trt.int32, trt.int64): + raise RuntimeError( + "The `embedding` op requires int32 or int64 indices, " + f"but received {indices_tensor.dtype}." + ) # unsupported parameters # ignore padding_idx, scale_grad_by_freq, and sparse # since they are meaningful for training only diff --git a/tests/py/dynamo/conversion/test_embedding_aten.py b/tests/py/dynamo/conversion/test_embedding_aten.py index c04d89ff9e..a9122e8ab0 100644 --- a/tests/py/dynamo/conversion/test_embedding_aten.py +++ b/tests/py/dynamo/conversion/test_embedding_aten.py @@ -29,6 +29,18 @@ class TestEmbeddingConverter(DispatchTestCase): weights_tensor=torch.randn((5, 10), dtype=torch.float32), sparse=True, ), + param( + test_name="1d_indices_int64", + indices_tensor=torch.tensor([3, 1, 2], dtype=torch.int64), + weights_tensor=torch.randn((5, 10), dtype=torch.float32), + sparse=False, + ), + param( + test_name="2d_indices_int64", + indices_tensor=torch.tensor([[3, 1, 2], [4, 1, 3]], dtype=torch.int64), + weights_tensor=torch.randn((5, 10), dtype=torch.float32), + sparse=True, + ), ] ) def test_embedding( @@ -57,8 +69,40 @@ def forward(self, indices, weights): inputs=[indices_tensor, weights_tensor], ) + def test_embedding_with_constant_int64_indices(self): + class ConstantInt64IndicesEmbedding(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer( + "indices", + torch.tensor([[0, 3, 1]], dtype=torch.int64), + persistent=False, + ) + + def forward(self, weights): + return torch.ops.aten.embedding.default( + weights, + self.indices, + -1, + False, + False, + ) + + self.run_test( + ConstantInt64IndicesEmbedding(), + inputs=[torch.randn((5, 10), dtype=torch.float32)], + ) + + @parameterized.expand( + [ + param(test_name="int32_indices", indices_dtype=torch.int32), + param(test_name="int64_indices", indices_dtype=torch.int64), + ] + ) def test_embedding_with_dynamic_shape_four_dimensions( self, + test_name, + indices_dtype, padding_idx=-1, max_norm=None, norm_type=2.0, @@ -78,7 +122,7 @@ def forward(self, input, weights): input_specs = [ Input( shape=(-1, -1, -1, -1), - dtype=torch.int32, + dtype=indices_dtype, shape_ranges=[((1, 1, 1, 1), (2, 3, 4, 5), (2, 3, 10, 10))], ), Input(