add is_binary=True (#3303)

pyro-ppl · Dec 6, 2023 · 9c4f932 · 9c4f932
1 parent a14fabc
commit 9c4f932
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pyro/distributions/zero_inflated.py b/pyro/distributions/zero_inflated.py
@@ -60,11 +60,11 @@ def support(self):
 
     @lazy_property
     def gate(self):
-        return logits_to_probs(self.gate_logits)
+        return logits_to_probs(self.gate_logits, is_binary=True)
 
     @lazy_property
     def gate_logits(self):
-        return probs_to_logits(self.gate)
+        return probs_to_logits(self.gate, is_binary=True)
 
     def log_prob(self, value):
         if self._validate_args: