diff --git a/dgl_sparse/src/elemenwise_op.cc b/dgl_sparse/src/elemenwise_op.cc index e25ca790f68a..616b395ab238 100644 --- a/dgl_sparse/src/elemenwise_op.cc +++ b/dgl_sparse/src/elemenwise_op.cc @@ -22,9 +22,15 @@ c10::intrusive_ptr SpSpAdd( const c10::intrusive_ptr& A, const c10::intrusive_ptr& B) { ElementwiseOpSanityCheck(A, B); - auto torch_A = COOToTorchCOO(A->COOPtr(), A->value()); - auto torch_B = COOToTorchCOO(B->COOPtr(), B->value()); - auto sum = (torch_A + torch_B).coalesce(); + torch::Tensor sum; + { + // TODO(#5145) This is a workaround to reduce peak memory usage. It is no + // longer needed after we address #5145. + auto torch_A = COOToTorchCOO(A->COOPtr(), A->value()); + auto torch_B = COOToTorchCOO(B->COOPtr(), B->value()); + sum = torch_A + torch_B; + } + sum = sum.coalesce(); auto indices = sum.indices(); auto row = indices[0]; auto col = indices[1];