search

tensorly / torch

TensorLy-Torch: Deep Tensor Learning with TensorLy and PyTorch
http://tensorly.org/torch/
BSD 3-Clause "New" or "Revised" License
70 stars 18 forks source link

BlockTT does not support CUDA Tensor as indices #23

Closed JeremieMelo closed 1 year ago

JeremieMelo commented 2 years ago

When I use FactorizedEmbedding, I found that it does not support CUDA Tensor as input indices. Because in BlockTT.__getitem__(), it uses np.unravel_index(index), which cannot take CUDA Tensor as input... Since PyTorch does not have unravel_index() yet, either moving CUDA indices to CPU or writing a Torch Version of unravel_index() will fix this bug.

JeanKossaifi commented 2 years ago

Do you have a usecase where you need the indices to be on GPU?

Typically, given the size of the indices, the overhead of converting is negligible, but if you have need for supporting torch tensors there we can add that.

JeanKossaifi commented 1 year ago

Fixed in #24