Made TorchLens compatible with sparse tensors.

johnmarktaylor91 · Jan 13, 2025 · 6745a86 · 6745a86
1 parent e07e784
commit 6745a86
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git 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
@@ -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):