Fixed issue #32, where some torchvision ops weren't being decorated (…

…since they're added separately to the PyTorch namespace)
johnmarktaylor91 · Dec 27, 2024 · 334c673 · 334c673
1 parent f38d342
commit 334c673
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -14,7 +14,7 @@
 
 setup(
     name="torchlens",
-    version="0.1.25",
+    version="0.1.26",
     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/constants.py b/torchlens/constants.py
@@ -415,5 +415,20 @@ def my_get_overridable_functions() -> List:
     return func_names
 
 
+TORCHVISION_FUNCS = [
+    ("torch.ops.torchvision.nms", "_op"),
+    ("torch.ops.torchvision.deform_conv2d", "_op"),
+    ("torch.ops.torchvision.ps_roi_align", "_op"),
+    ("torch.ops.torchvision.ps_roi_pool", "_op"),
+    ("torch.ops.torchvision.roi_align", "_op"),
+    ("torch.ops.torchvision.roi_pool", "_op")]
+
 OVERRIDABLE_FUNCS = my_get_overridable_functions()
 ORIG_TORCH_FUNCS = OVERRIDABLE_FUNCS + IGNORED_FUNCS
+
+try:
+    import torchvision
+
+    ORIG_TORCH_FUNCS += TORCHVISION_FUNCS
+except ModuleNotFoundError:
+    pass