Skip to content

Commit b2ba34b

Browse files
committed
bw fix
1 parent 1a4bdd3 commit b2ba34b

File tree

1 file changed

+8
-6
lines changed

1 file changed

+8
-6
lines changed

torch_sparse/spspmm.py

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -81,11 +81,13 @@ def from_scipy(A):
8181

8282

8383
def lift(indexA, valueA, indexB, n): # pragma: no cover
84-
indexA = indexA[0] * n + indexA[1]
85-
indexB = indexB[0] * n + indexB[1]
84+
idxA = indexA[0] * n + indexA[1]
85+
idxB = indexB[0] * n + indexB[1]
8686

87-
value = valueA.new_zeros(indexB.max().item() + 1)
88-
value[indexA] = valueA
89-
value = value[indexB]
87+
max_value = max(idxA.max().item(), idxB.max().item()) + 1
88+
valueB = valueA.new_zeros(max_value)
9089

91-
return value
90+
valueB[idxA] = valueA
91+
valueB = valueB[idxB]
92+
93+
return valueB

0 commit comments

Comments
 (0)