fix: added a simple way to add additional handlers

MLRichter · Jan 28, 2022 · 8f2cefd · 8f2cefd
1 parent d968884
commit 8f2cefd
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/rfa_toolbox/encodings/pytorch/__init__.py b/rfa_toolbox/encodings/pytorch/__init__.py
@@ -1,2 +1,5 @@
 # flake8: noqa: F401
-from rfa_toolbox.encodings.pytorch.utils import toggle_coerce_torch_functional
+from rfa_toolbox.encodings.pytorch.utils import (
+    add_custom_layer_handler,
+    toggle_coerce_torch_functional,
+)
diff --git a/tests/test_encodings/test_pytorch.py b/tests/test_encodings/test_pytorch.py
@@ -108,7 +108,7 @@ def test_adding_a_handler(self):
         )
 
         from rfa_toolbox import create_graph_from_pytorch_model
-        from rfa_toolbox.encodings.pytorch.utils import add_custom_layer_handler
+        from rfa_toolbox.encodings.pytorch import add_custom_layer_handler
 
         def kernel_size(module):
             x1 = module.k_size