TensorRT icon indicating copy to clipboard operation
TensorRT copied to clipboard

🐛 [Bug] arange converter fails with ITensor input

Open peri044 opened this issue 1 year ago • 0 comments

Bug Description

arange converter fails with ITensor input

  File "/home/dperi/Downloads/TensorRT/py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py", line 367, in call_function
    return converter(self.ctx, target, args, kwargs, self._cur_node_name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/dperi/Downloads/TensorRT/py/torch_tensorrt/dynamo/conversion/ops_evaluators.py", line 49, in aten_ops_arange_start_step
    return np.arange(*args)
           ^^^^^^^^^^^^^^^^
TypeError: unsupported operand type(s) for -: 'tensorrt.tensorrt.ITensor' and 'int'

While executing %arange : [num_users=1] = call_function[target=torch.ops.aten.arange.start_step](args = (0, %add), kwargs = {dtype: torch.int64, layout: torch.strided, device: cuda:0, pin_memory: False, _itensor_to_tensor_meta: {<tensorrt.tensorrt.ITensor object at 0x7f35ba688d30>: ((1, 7), torch.int64, False, (7, 1), torch.contiguous_format, False, {}), <tensorrt.tensorrt.ITensor object at 0x7f35ba689130>: None, <tensorrt.tensorrt.ITensor object at 0x7f35ba689330>: None}})

To Reproduce

Install transformers via pip install transformers Run the following script

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import torch_tensorrt

torch_device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained("gpt2")

# add the EOS token as PAD token to avoid warnings
model = AutoModelForCausalLM.from_pretrained("gpt2", pad_token_id=tokenizer.eos_token_id).to(torch_device)
model_inputs = tokenizer('I enjoy walking with my cute dog', return_tensors='pt').to(torch_device)

dyn_sl = torch.export.Dim("seq_length", max=1023)
ep = torch.export.export(model, (model_inputs['input_ids'],), dynamic_shapes=({1: dyn_sl},))
trt_ep = torch_tensorrt.dynamo.compile(ep, model_inputs['input_ids'], debug=True, truncate_long_and_double=True)

Expected behavior

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

  • Torch-TensorRT Version (e.g. 1.0.0):
  • PyTorch Version (e.g. 1.0):
  • CPU Architecture:
  • OS (e.g., Linux):
  • How you installed PyTorch (conda, pip, libtorch, source):
  • Build command you used (if compiling from source):
  • Are you using local sources or building from archives:
  • Python version:
  • CUDA version:
  • GPU models and configuration:
  • Any other relevant information:

Additional context

peri044 avatar Feb 01 '24 17:02 peri044