https://github.com/nv-tlabs/LION/issues/31#issue-1627736930 the original code index cpu tensor with cuda tensor (works in torch 1.10.2), may fail in other torch version?