Fixed small bug that happens when parameters are passed into a module…

…'s forward pass.
johnmarktaylor91 · Sep 24, 2023 · 51ab52a · 51ab52a
1 parent 3556f9b
commit 51ab52a
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -16,7 +16,7 @@
 
 setup(
     name="torchlens",
-    version="0.1.8",
+    version="0.1.9",
     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
@@ -1294,7 +1294,7 @@ def decorated_forward(*args, **kwargs):
             module_pass_label = (module_address, module.tl_module_pass_num)
             module.tl_module_pass_labels.append(module_pass_label)
             input_tensors = get_vars_of_type_from_obj(
-                [args, kwargs], torch.Tensor, search_depth=4
+                [args, kwargs], torch.Tensor, [torch.nn.Parameter], search_depth=4
             )
             input_tensor_labels = set()
             for t in input_tensors: