Add End2EndDomainModule to __init__

ruflab · Sep 19, 2024 · 9976918 · 9976918
1 parent 64e69a1
commit 9976918
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/shimmer/__init__.py b/shimmer/__init__.py
@@ -10,7 +10,7 @@
     ContrastiveLossType,
     contrastive_loss,
 )
-from shimmer.modules.domain import DomainModule, LossOutput
+from shimmer.modules.domain import DomainModule, End2EndDomainModule, LossOutput
 from shimmer.modules.global_workspace import (
     GlobalWorkspace2Domains,
     GlobalWorkspaceBase,
@@ -79,6 +79,7 @@
     "pretrained_global_workspace",
     "LossOutput",
     "DomainModule",
+    "End2EndDomainModule",
     "GWDecoder",
     "GWEncoder",
     "GWEncoderLinear",

diff --git a/shimmer/modules/__init__.py b/shimmer/modules/__init__.py
@@ -5,7 +5,7 @@
     ContrastiveLossType,
     contrastive_loss,
 )
-from shimmer.modules.domain import DomainModule, LossOutput
+from shimmer.modules.domain import DomainModule, End2EndDomainModule, LossOutput
 from shimmer.modules.global_workspace import (
     GlobalWorkspace2Domains,
     GlobalWorkspaceBase,
@@ -59,6 +59,7 @@
     "pretrained_global_workspace",
     "LossOutput",
     "DomainModule",
+    "End2EndDomainModule",
     "GWDecoder",
     "GWEncoder",
     "GWEncoderLinear",