The input and src tensor can be on cpu but the index tensor is being casted to gpu:0 and this results in src and index tensor are on different devices issue. Cast it to the same device as input.
To Reproduce
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:
Bug Description
The input and src tensor can be on cpu but the index tensor is being casted to gpu:0 and this results in src and index tensor are on different devices issue. Cast it to the same device as input.
To Reproduce
Expected behavior
Environment
conda
,pip
,libtorch
, source):Additional context