use randn32

FluxML · Jul 12, 2021 · dfb390d · dfb390d
1 parent 058a4a0
commit dfb390d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/src/layers/basic.jl b/src/layers/basic.jl
@@ -471,11 +471,8 @@ end
 
 @functor Embedding
 
-function Embedding(in::Integer, out::Integer;
-               init = (i...) -> randn(Float32, i...))
-  return Embedding(init(out, in))
-end
-
+Embedding(in::Integer, out::Integer; init = randn32) = Embedding(init(out, in))
+
 (m::Embedding)(x::Union{OneHotVector, OneHotMatrix}) = m.weight * x # equivalent to m.weight[:,onecold(x)]
 (m::Embedding)(x::Integer) = m([x])
 (m::Embedding)(x::AbstractVector) = NNlib.gather(m.weight, x)

diff --git a/src/utils.jl b/src/utils.jl
@@ -368,6 +368,8 @@ identity_init(rng::AbstractRNG; init_kwargs...) = (args...;kwargs...) -> identit
 
 ones32(dims...) = Base.ones(Float32, dims...)
 zeros32(dims...) = Base.zeros(Float32, dims...)
+rand32(dims...) = Base.rand(Float32, dims...)
+randn32(dims...) = Base.randn(Float32, dims...)
 
 """
     create_bias(weights, bias, length)