relax bernoulli rand type

xukai92 · Apr 23, 2019 · f9d7a23 · f9d7a23
1 parent 5bfb9c4
commit f9d7a23
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/distributions/bernoulli.jl b/src/distributions/bernoulli.jl
@@ -17,11 +17,11 @@ function logpdf(bb::BatchBernoulli, x)
     return lp
 end
 
-function rand(bb::BatchBernoulli)
+function rand(bb::BatchBernoulli{T}) where {T}
     FT = eltype(bb.p)
 
     u = rand(FT, size(bb.p)...)
-    x = AT((Array(bb.p) .> u) * one(FT))
+    x = T((Array(bb.p) .> u) * one(FT))
     return x
 end