Fix dtype index bug.

facebookresearch · Dec 5, 2022 · 4a78a03 · 4a78a03
1 parent 28778a5
commit 4a78a03
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/theseus/utils/sparse_matrix_utils.py b/theseus/utils/sparse_matrix_utils.py
@@ -130,7 +130,7 @@ def _sparse_mat_vec_bwd_backend(
     for row in range(num_rows):
         start = A_row_ptr[row]
         end = A_row_ptr[row + 1]
-        columns = A_col_ind[start:end]
+        columns = A_col_ind[start:end].long()
         if is_tmat:
             A_grad[:, start:end] = v[:, row].view(-1, 1) * grad_output[:, columns]
             v_grad[:, row] = (grad_output[:, columns] * A_val[:, start:end]).sum(dim=1)