From c171f37204b2d82b58849547fc80fffb0983843a Mon Sep 17 00:00:00 2001
From: DangTAnh <163684532+DangTAnh@users.noreply.github.com>
Date: Thu, 8 Oct 2026 09:11:27 +0700
Subject: [PATCH 1/2] fix(summary): support sparse tensor inputs in memory
 accounting

Sparse COO/CSR/CSC adjacency inputs (e.g. GCNs) crashed get_total_memory_used via untyped_storage, so no summary was produced. Sum member-tensor storage for sparse layouts instead.

Fixes #150
---
 tests/torchinfo_test.py | 30 ++++++++++++++++++++++++++++++
 torchinfo/torchinfo.py  | 40 +++++++++++++++++++++++++++++++++++-----
 2 files changed, 65 insertions(+), 5 deletions(-)

diff --git a/tests/torchinfo_test.py b/tests/torchinfo_test.py
index e91ff82..562ecda 100644
--- a/tests/torchinfo_test.py
+++ b/tests/torchinfo_test.py
@@ -467,6 +467,36 @@ def test_multiple_input_tensor_list() -> None:
     assert metrics.input_size == [torch.Size([1, 300]), torch.Size([1, 300])]
 
 
+def test_sparse_input() -> None:
+    # https://github.com/TylerYep/torchinfo/issues/150
+    class GraphConvolution(nn.Module):
+        def __init__(self, in_features: int, out_features: int) -> None:
+            super().__init__()
+            self.weight = nn.Parameter(torch.randn(in_features, out_features))
+
+        def forward(self, x: torch.Tensor, adj: torch.Tensor) -> torch.Tensor:
+            return torch.spmm(adj, x @ self.weight)  # type: ignore[no-any-return]
+
+    class GCN(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.gc1 = GraphConvolution(16, 8)
+            self.gc2 = GraphConvolution(8, 4)
+
+        def forward(self, x: torch.Tensor, adj: torch.Tensor) -> torch.Tensor:
+            return self.gc2(torch.relu(self.gc1(x, adj)), adj)
+
+    features = torch.randn(10, 16)
+    adj = torch.sparse_coo_tensor(
+        torch.tensor([[0, 1, 2], [1, 2, 0]]), torch.ones(3), (10, 10)
+    ).coalesce()
+
+    metrics = summary(GCN(), input_data=[features, adj], verbose=0)
+
+    assert metrics.total_params == 16 * 8 + 8 * 4 == 160
+    assert metrics.trainable_params == 160
+
+
 def test_namedtuple() -> None:
     model = NamedTuple()
     input_size = [(2, 1, 28, 28), (2, 1, 28, 28)]
diff --git a/torchinfo/torchinfo.py b/torchinfo/torchinfo.py
index 46b4502..682d0fb 100644
--- a/torchinfo/torchinfo.py
+++ b/torchinfo/torchinfo.py
@@ -507,15 +507,45 @@ def get_input_data_sizes(data: Any) -> Any:
     )
 
 
+def get_tensor_memory(data: torch.Tensor) -> int:
+    """Calculates the memory of a single tensor, including sparse layouts."""
+    try:
+        storage = (
+            data.untyped_storage()
+            if hasattr(data, "untyped_storage")
+            else data.storage()  # type: ignore[no-untyped-call]
+        )
+        return sys.getsizeof(storage)
+    except NotImplementedError:
+        # Sparse tensors (COO, CSR, ...) expose no single storage,
+        # so sum the member tensors instead. Accessors differ per
+        # layout, so probe the known ones and keep what succeeds.
+        members: list[torch.Tensor] = [data.values()]
+        for name in (
+            "indices",
+            "crow_indices",
+            "col_indices",
+            "ccol_indices",
+            "row_indices",
+        ):
+            getter = getattr(data, name, None)
+            if getter is None:
+                continue
+            try:
+                member = getter()
+            except NotImplementedError:
+                continue
+            except RuntimeError:
+                continue
+            members.append(member)
+        return sum(get_tensor_memory(member) for member in members)
+
+
 def get_total_memory_used(data: CORRECTED_INPUT_DATA_TYPE) -> int:
     """Calculates the total memory of all tensors stored in data."""
     result = traverse_input_data(
         data,
-        action_fn=lambda data: sys.getsizeof(
-            data.untyped_storage()
-            if hasattr(data, "untyped_storage")
-            else data.storage()
-        ),
+        action_fn=get_tensor_memory,
         aggregate_fn=(
             # We don't need the dictionary keys in this case
             lambda data: (

From 2f9a5b84c317c5e4a1660dfd320f5c189bc0797f Mon Sep 17 00:00:00 2001
From: DangTAnh <163684532+DangTAnh@users.noreply.github.com>
Date: Fri, 9 Oct 2026 19:15:08 +0700
Subject: [PATCH 2/2] fix(tests): silence no-any-return via annotated Tensor in
 sparse test

---
 tests/torchinfo_test.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tests/torchinfo_test.py b/tests/torchinfo_test.py
index 562ecda..73a9477 100644
--- a/tests/torchinfo_test.py
+++ b/tests/torchinfo_test.py
@@ -475,7 +475,7 @@ def __init__(self, in_features: int, out_features: int) -> None:
             self.weight = nn.Parameter(torch.randn(in_features, out_features))
 
         def forward(self, x: torch.Tensor, adj: torch.Tensor) -> torch.Tensor:
-            return torch.spmm(adj, x @ self.weight)  # type: ignore[no-any-return]
+            return torch.spmm(adj, x @ self.weight)
 
     class GCN(nn.Module):
         def __init__(self) -> None:
@@ -484,7 +484,8 @@ def __init__(self) -> None:
             self.gc2 = GraphConvolution(8, 4)
 
         def forward(self, x: torch.Tensor, adj: torch.Tensor) -> torch.Tensor:
-            return self.gc2(torch.relu(self.gc1(x, adj)), adj)
+            out: torch.Tensor = self.gc2(torch.relu(self.gc1(x, adj)), adj)
+            return out
 
     features = torch.randn(10, 16)
     adj = torch.sparse_coo_tensor(