Fix type-stability of optimizer

Author: pxl-th

Project.toml

@@ -1,7 +1,7 @@
 name = "NerfUtils"
 uuid = "99c1d5ce-7c61-4a25-a107-a5ade2e2a8e4"
 authors = ["Anton Smirnov <tonysmn97@gmail.com>"]
-version = "0.1.3"
+version = "0.1.4"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/nn/adam.jl

@@ -1,3 +1,7 @@
+_eltype(θ::T) where T <: Union{Tuple, NamedTuple} = _eltype(first(θ))
+_eltype(θ::T) where T <: AbstractVector = T
+_eltype(θ::T) where T = _eltype(reshape(θ, :))
+
 """
     Adam(kab, θ; kwargs...)
 
@@ -20,7 +24,8 @@ end
 KernelAbstractions.get_backend(opt::Adam) = get_backend(first(opt.μ))
 
 function Adam(kab, θ; kwargs...)
-    μ, ν = [], [] # FIXME unstable
+    T = _eltype(θ)
+    μ, ν = T[], T[]
     _add_moments!(μ, ν, θ, kab)
     Adam(; μ, ν, kwargs...)
 end