TensorRT
TensorRT copied to clipboard
🐛 [Bug] arange converter fails with ITensor input
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: