[Feature Request] Add one dimensional tensor support for `add_trt_constant`

NVIDIA-AI-IOT / torch2trt

An easy to use PyTorch to TensorRT converter

MIT License

4.58k stars 675 forks source link

[Feature Request] Add one dimensional tensor support for `add_trt_constant` #765

Open chaoz-dev opened 2 years ago

chaoz-dev commented 2 years ago

Currently, add_trt_constant assumes dim > 1 for input tensors (where the 1st dimension is probably assumed to be the batch dimension).

However, there are situations where we may use one dimensional tensors as input to add_trt_constant, especially now that explicit batch dimensions are supported.

We should add support for the one dimensional use case, as the current implementation with tensor[0] will return the first element of the tensor, which is not the desired behavior.

chaoz-dev commented 2 years ago

I'll put up a PR up for this.

chaoz-dev commented 2 years ago

This is needed for #755