Add latent_dim to abstract GWModule for auto-completion

ruflab · Dec 14, 2023 · 663a525 · 663a525
1 parent 6e8095f
commit 663a525
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/shimmer/modules/gw_module.py b/shimmer/modules/gw_module.py
@@ -80,6 +80,7 @@ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
 
 class GWModule(nn.Module):
     domain_descr: Mapping[str, DomainDescription]
+    latent_dim: int
 
     def fusion_mechanism(self, x: Mapping[str, torch.Tensor]) -> torch.Tensor:
         raise NotImplementedError