update setparams_varinfo!! for MH

Author: penelopeysm

src/mcmc/gibbs.jl

@@ -540,14 +540,11 @@ function setparams_varinfo!!(model, ::Sampler, state, params::AbstractVarInfo)
 end
 
 function setparams_varinfo!!(
-    model::DynamicPPL.Model,
-    sampler::Sampler{<:MH},
-    state::AbstractVarInfo,
-    params::AbstractVarInfo,
+    model::DynamicPPL.Model, sampler::Sampler{<:MH}, state::MHState, params::AbstractVarInfo
 )
-    # The state is already a VarInfo, so we can just return `params`, but first we need to
-    # update its logprob.
-    return last(DynamicPPL.evaluate!!(model, params))
+    # Re-evaluate to update the logprob.
+    new_vi = last(DynamicPPL.evaluate!!(model, params))
+    return MHState(new_vi, DynamicPPL.getlogjoint_internal(new_vi))
 end
 
 function setparams_varinfo!!(