remove learning logit scale in deterministic mode

ruflab · Nov 8, 2023 · e118b13 · e118b13
1 parent a91f05a
commit e118b13
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/shimmer/modules/losses.py b/shimmer/modules/losses.py
@@ -288,7 +288,7 @@ def __init__(
         self.gw_mod = gw_mod
         self.domain_mods = domain_mods
         self.loss_coefs = coef_buffers
-        self.logit_scale = torch.nn.Parameter(torch.tensor([1 / 0.07]).log())
+        self.register_buffer("logit_scale", torch.tensor([1 / 0.01]).log())
 
     def demi_cycle_loss(
         self, latent_domains: LatentsT