diff --git a/pyCP_APR/torch_backend/ttv_sptensor.py b/pyCP_APR/torch_backend/ttv_sptensor.py index a31cc8f..5062325 100644 --- a/pyCP_APR/torch_backend/ttv_sptensor.py +++ b/pyCP_APR/torch_backend/ttv_sptensor.py @@ -26,8 +26,8 @@ def ttv(M, vecs, dims=[]): product of KRUSKAL tensor X with a (column) vector vecs. """ - dims = tr.arange(M.Dimensions) - vidx = tr.arange(M.Dimensions) + dims = tr.arange(M.Dimensions).to(M.device) + vidx = tr.arange(M.Dimensions).to(M.device) combined = tr.cat((dims, vidx)) uniques, counts = combined.unique(return_counts=True) diff --git a/setup.py b/setup.py index c292d97..c9525df 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ description='pyCP_APR: CP-APR Tensor Decomposition with PyTorch Backend.', long_description=LONG_DESCRIPTION, long_description_content_type='text/markdown', - package_dir={'pyCP_APR': 'pyCP_APR/'}, + package_dir={'pyCP_APR': 'pyCP_APR'}, platforms = ["Linux", "Mac", "Windows"], include_package_data=True, setup_requires=[