fixing test

Author: zuhengxu

src/optimize.jl

@@ -66,9 +66,9 @@ Iteratively updating the parameters `θ` of the normalizing flow `re(θ)` by cal
 """
 function optimize(
     adbackend,
-    loss::Function, 
+    loss,
     θ₀::AbstractVector{<:Real}, 
-    reconstruct::Function,
+    reconstruct,
     args...;
     max_iters::Int=10000,
     optimiser::Optimisers.AbstractRule=Optimisers.ADAM(),

test/interface.jl

@@ -13,6 +13,7 @@
             target = MvNormal(μ, Σ)
             logp(z) = logpdf(target, z)
 
+            @leaf MvNormal
             q₀ = MvNormal(zeros(T, 2), ones(T, 2))
             flow = Bijectors.transformed(
                 q₀, Bijectors.Shift(zero.(μ)) ∘ Bijectors.Scale(ones(T, 2))
@@ -21,7 +22,7 @@
             sample_per_iter = 10
             cb(iter, opt_stats, re, θ) = (sample_per_iter=sample_per_iter,)
             checkconv(iter, stat, re, θ, st) = stat.gradient_norm < 1e-3
-            flow_trained, stats, _ = train_flow(
+            flow_trained, stats, _, _ = train_flow(
                 elbo,
                 flow,
                 logp,