add logit scale def for lsps..

bdvllrs · bdvllrs · commit ae4a2385d193 · 2023-11-08T13:43:55.000Z
diff --git a/shimmer/modules/losses.py b/shimmer/modules/losses.py
@@ -278,6 +278,8 @@ def _contrastive_loss_with_uncertainty(
 
 
 class DeterministicGWLosses(GWLosses):
+    logit_scale: torch.Tensor
+
     def __init__(
         self,
         gw_mod: DeterministicGWModule,
