Reset theta after find init eps

TuringLang · Jun 19, 2017 · 6b78aa4 · 6b78aa4
1 parent 5e1be20
commit 6b78aa4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/src/samplers/hmcda.jl b/src/samplers/hmcda.jl
@@ -54,9 +54,11 @@ function step(model, spl::Sampler{HMCDA}, vi::VarInfo, is_first::Bool)
 
       init_warm_up_params(vi, spl)
 
+      oldθ = vi[spl]
       ϵ = spl.alg.delta > 0 ?
           find_good_eps(model, vi, spl) :       # heuristically find optimal ϵ
           spl.info[:pre_set_ϵ]
+      vi[spl] = oldθ
 
       if spl.alg.gid != 0 invlink!(vi, spl) end # R -> X
 

diff --git a/src/samplers/nuts.jl b/src/samplers/nuts.jl
@@ -49,7 +49,9 @@ function step(model::Function, spl::Sampler{NUTS}, vi::VarInfo, is_first::Bool)
 
       init_warm_up_params(vi, spl)
 
+      oldθ = vi[spl]
       ϵ = find_good_eps(model, vi, spl)           # heuristically find optimal ϵ
+      vi[spl] = oldθ
 
       if spl.alg.gid != 0 invlink!(vi, spl) end   # R -> X