wrong device in ttv_tensor.py bug fixed

lanl · Dec 18, 2023 · 5c42f6c · 5c42f6c
1 parent bdaed6e
commit 5c42f6c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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)