fix(cat): correct temperature scaling

braun-steven · Nov 6, 2024 · 8cabee2 · 8cabee2
1 parent a3f93f5
commit 8cabee2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/simple_einet/layers/distributions/abstract_leaf.py b/simple_einet/layers/distributions/abstract_leaf.py
@@ -124,7 +124,7 @@ def dist_sample(distribution: dist.Distribution, ctx: SamplingContext = None) ->
             elif type(distribution) == CustomNormal:
                 distribution = CustomNormal(mu=distribution.mu, sigma=distribution.sigma * np.sqrt(ctx.temperature_leaves))
             elif type(distribution) == dist.Categorical:
-                distribution = dist.Categorical(logits=F.softmax(distribution.logits / ctx.temperature_leaves))
+                distribution = dist.Categorical(logits=distribution.logits / ctx.temperature_leaves)
 
             samples = distribution.sample(sample_shape=(ctx.num_samples,)).float()