diff --git a/src/optimize_flow.jl b/src/optimize_flow.jl
index 7ccebd8..063efd5 100644
--- a/src/optimize_flow.jl
+++ b/src/optimize_flow.jl
@@ -12,28 +12,29 @@ function negll_flow_loss(flow::F, x::AbstractMatrix{<:Real}, logd_orig::Abstract
 end
 
 function negll_flow(flow::F, x::AbstractMatrix{<:Real}, logd_orig::AbstractVector, logpdf::Tuple{Function, Function}) where F<:AbstractFlow
-    negll, back = Zygote.pullback(negll_flow, flow, x, logd_orig, logpdf[2])
+    negll, back = Zygote.pullback(negll_flow_loss, flow, x, logd_orig, logpdf[2])
     d_flow = back(one(eltype(x)))[1]
     return negll, d_flow
 end
 export negll_flow
 
-function KLDiv_flow_loss(flow::F, x::AbstractMatrix{<:Real}, logd_orig::AbstractVector, logpdfs::Tuple{Function, Function}) where F<:AbstractFlow
+function KLDiv_flow_loss(flow::F, x::AbstractMatrix{<:Real}, logd_orig::AbstractVector, logpdf::Function) where F<:AbstractFlow
     nsamples = size(x, 2) 
-    flow_corr = fchain(flow,logpdfs[2].f)
-    logpdf_y = logpdfs[2].logdensity
+    flow_corr = fchain(flow,logpdf.f)
+    #logpdf_y = logpdfs[2].logdensity
     y, ladj = with_logabsdet_jacobian(flow_corr, x)
-    KLDiv = sum(exp.(logd_orig - vec(ladj)) .* (logd_orig - vec(ladj) - logpdf_y(y))) / nsamples
+    KLDiv = sum(exp.(logd_orig - vec(ladj)) .* (logd_orig - vec(ladj) - logpdf.logdensity(y))) / nsamples
     return KLDiv
 end
 
-function KLDiv_flow(flow::F, x::AbstractMatrix{<:Real}, logd_orig::AbstractVector, logpdfs::Tuple{Function, Function}) where F<:AbstractFlow
-    KLDiv, back = Zygote.pullback(KLDiv_flow_loss, flow, x, logd_orig, logpdfs)
+function KLDiv_flow(flow::F, x::AbstractMatrix{<:Real}, logd_orig::AbstractVector, logpdf::Tuple{Function, Function}) where F<:AbstractFlow
+    KLDiv, back = Zygote.pullback(KLDiv_flow_loss, flow, x, logd_orig, logpdf[2])
     d_flow = back(one(eltype(x)))[1]
     return KLDiv, d_flow
 end
 export KLDiv_flow
 
+
 function optimize_flow(samples::Union{Matrix, Tuple{Matrix, Matrix}}, 
     initial_flow::F where F<:AbstractFlow, 
     optimizer;