make sure kwargs pass to callback

TuringLang · Apr 10, 2021 · 0a8ae72 · 0a8ae72
1 parent a182954
commit 0a8ae72
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/sample.jl b/src/sample.jl
@@ -197,7 +197,7 @@ function mcmcsample(
         end
 
         # Run callback.
-        callback === nothing || callback(rng, model, sampler, sample, state, 1)
+        callback === nothing || callback(rng, model, sampler, sample, state, 1; kwargs...)
 
         # Save the sample.
         samples = AbstractMCMC.samples(sample, model, sampler; kwargs...)
@@ -217,7 +217,7 @@ function mcmcsample(
             sample, state = step(rng, model, sampler, state; kwargs...)
 
             # Run callback.
-            callback === nothing || callback(rng, model, sampler, sample, state, i)
+            callback === nothing || callback(rng, model, sampler, sample, state, i; kwargs...)
 
             # Save the sample.
             samples = save!!(samples, sample, i, model, sampler; kwargs...)