Fix wrong softmax projection on probs/logits (John Leland)

braun-steven · Jun 17, 2024 · 29fb588 · 29fb588
1 parent 0213fe0
commit 29fb588
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
@@ -112,7 +112,7 @@ def dist_sample(distribution: dist.Distribution, ctx: SamplingContext = None) ->
         elif type(distribution) == CustomNormal:
             distribution = CustomNormal(mu=distribution.mu, sigma=distribution.sigma / ctx.temperature_leaves)
         elif type(distribution) == dist.Categorical:
-            distribution = dist.Categorical(logits=F.log_softmax(distribution.probs / ctx.temperature_leaves))
+            distribution = dist.Categorical(logits=F.log_softmax(distribution.logits / ctx.temperature_leaves))
         samples = distribution.sample(sample_shape=(ctx.num_samples,)).float()
 
     assert (