Do not save unused data tensor for SDD operation.

stanford-futuredata · Aug 15, 2023 · 8aa616a · 8aa616a
1 parent d5c2dc6
commit 8aa616a
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/stk/backend/sputnik.py b/stk/backend/sputnik.py
@@ -247,7 +247,6 @@ def forward(ctx,
         ctx.save_for_backward(
             lhs,
             rhs,
-            data,
             offsets,
             row_indices,
             column_indices,
@@ -257,8 +256,8 @@ def forward(ctx,
         ctx.shape = shape
         out = torch.empty(
             data.shape,
-            dtype=data.dtype,
-            device=data.device)
+            dtype=lhs.dtype,
+            device=lhs.device)
         backend.sdd(lhs,
                     rhs,
                     shape,
@@ -272,7 +271,7 @@ def forward(ctx,
     @custom_bwd
     def backward(ctx, dy):
         lhs, rhs = ctx.saved_tensors[:2]
-        dy = (ctx.shape, dy) + ctx.saved_tensors[3:]
+        dy = (ctx.shape, dy) + ctx.saved_tensors[2:]
         trans_a = _is_transposed(lhs)
         trans_b = _is_transposed(rhs)