add T.clip to bernoulli distribution

haowen-xu · Dec 10, 2021 · 6354a95 · 6354a95
1 parent c55ca34
commit 6354a95
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tensorkit/distributions/bernoulli.py b/tensorkit/distributions/bernoulli.py
@@ -100,7 +100,11 @@ def logits(self) -> T.Tensor:
     def probs(self) -> T.Tensor:
         """Get the probability of being 1."""
         if self._probs is None:
-            self._probs = T.random.bernoulli_logits_to_probs(self._logits)
+            self._probs = T.clip(
+                T.random.bernoulli_logits_to_probs(self._logits),
+                0.,
+                1.
+            )
         return self._probs
 
     def _sample(self,