From 6745a86145eeb7284086a32b868e3c236d891eb3 Mon Sep 17 00:00:00 2001 From: JohnMark Taylor Date: Mon, 13 Jan 2025 10:36:22 -0500 Subject: [PATCH] Made TorchLens compatible with sparse tensors. --- setup.py | 2 +- torchlens/helper_funcs.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index c1eeb1d..f433bf0 100644 --- a/setup.py +++ b/setup.py @@ -14,7 +14,7 @@ setup( name="torchlens", - version="0.1.26", + version="0.1.27", description="A package for extracting activations from PyTorch models", long_description="A package for extracting activations from PyTorch models. Contains functionality for " "extracting model activations, visualizing a model's computational graph, and " diff --git a/torchlens/helper_funcs.py b/torchlens/helper_funcs.py index 52fae65..ac153b4 100644 --- a/torchlens/helper_funcs.py +++ b/torchlens/helper_funcs.py @@ -597,7 +597,7 @@ def get_tensor_memory_amount(t: torch.Tensor) -> int: cpu_data = clean_cpu(t.data) if cpu_data.dtype == torch.bfloat16: cpu_data = clean_to(cpu_data, torch.float16) - return getsizeof(np.array(cpu_data)) + return getsizeof(np.array(clean_dense(cpu_data))) def human_readable_size(size: int, decimal_places: int = 1) -> str: @@ -627,6 +627,7 @@ def human_readable_size(size: int, decimal_places: int = 1) -> str: clean_cpu = copy.deepcopy(torch.Tensor.cpu) clean_cuda = copy.deepcopy(torch.Tensor.cuda) clean_to = copy.deepcopy(torch.Tensor.to) +clean_dense = copy.deepcopy(torch.Tensor.to_dense) def print_override(t: torch.Tensor, func_name: str):