fix(summary): support sparse tensor inputs in memory accounting

tests/torchinfo_test.py

@@ -467,6 +467,37 @@ 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)
+
+    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:
+            out: torch.Tensor = self.gc2(torch.relu(self.gc1(x, adj)), adj)
+            return out
+
+    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)]

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: (