Fix attribute error

ruflab · Oct 8, 2024 · 975a6e0 · 975a6e0
1 parent 6259db4
commit 975a6e0
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/shimmer/modules/gw_module.py b/shimmer/modules/gw_module.py
@@ -374,14 +374,17 @@ def __init__(
             fusion_activation_fn (`Callable[[torch.Tensor], torch.Tensor]`): activation
                 function used to fuse the domains.
         """
-        super().__init__(domain_modules, workspace_dim, fusion_activation_fn)
+        super().__init__(domain_modules, workspace_dim)
 
         self.gw_encoders = nn.ModuleDict(gw_encoders)
         """The module's encoders"""
 
         self.gw_decoders = nn.ModuleDict(gw_decoders)
         """The module's decoders"""
 
+        self.fusion_activation_fn = fusion_activation_fn
+        """Activation function used to fuse the domains."""
+
     def fuse(
         self,
         x: LatentsDomainGroupT,