Uncertainty GW loss uses GWFusionLoss (#60)

ruflab · Apr 16, 2024 · fea876f · fea876f
1 parent 7671609
commit fea876f
Show file tree

Hide file tree

Showing 3 changed files with 224 additions and 266 deletions.
diff --git a/shimmer/modules/global_workspace.py b/shimmer/modules/global_workspace.py
@@ -635,7 +635,7 @@ def __init__(
         gw_encoders: Mapping[str, Module],
         gw_decoders: Mapping[str, Module],
         workspace_dim: int,
-        loss_coefs: LossCoefs,
+        loss_coefs: BroadcastLossCoefs,
         selection_temperature: float = 0.2,
         optim_lr: float = 1e-3,
         optim_weight_decay: float = 0.0,