Allow using heads by importing pytorch_toolbelt.modules

BloodAxe · Jan 7, 2024 · ec571e8 · ec571e8
1 parent b22838f
commit ec571e8
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/pytorch_toolbelt/modules/__init__.py b/pytorch_toolbelt/modules/__init__.py
@@ -17,3 +17,5 @@
 from .decoders import *
 from .initialization import *
 from .normalization import *
+
+from .heads import *