Fixed issue #32, where some modules contain a "BufferList" and the bu…

…ffers aren't found.
johnmarktaylor91 · Dec 27, 2024 · f38d342 · f38d342
1 parent 844ed04
commit f38d342
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/setup.py b/setup.py
@@ -14,7 +14,7 @@
 
 setup(
     name="torchlens",
-    version="0.1.24",
+    version="0.1.25",
     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/model_history.py b/torchlens/model_history.py
@@ -1317,10 +1317,11 @@ def prepare_buffer_tensors(self, model: nn.Module):
         """
         submodules = self.get_all_submodules(model)
         for submodule in submodules:
-            for attribute_name, attribute in iter_accessible_attributes(submodule):
+            attr_list = list(submodule.named_buffers()) + list(iter_accessible_attributes(submodule))
+            for attribute_name, attribute in attr_list:
                 if issubclass(type(attribute), torch.Tensor) and not issubclass(
                         type(attribute), torch.nn.Parameter
-                ):
+                ) and not hasattr(attribute, 'tl_buffer_address'):
                     if submodule.tl_module_address == "":
                         buffer_address = attribute_name
                     else: