Only update NUTS divergence stats after tuning

pymc-devs · Jan 9, 2025 · 956ab12 · 956ab12
1 parent d336ad8
commit 956ab12
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/pymc/step_methods/hmc/nuts.py b/pymc/step_methods/hmc/nuts.py
@@ -253,7 +253,9 @@ def update_stats(stats, step_stats, chain_idx):
             if isinstance(step_stats, list):
                 step_stats = step_stats[0]
 
-            stats["divergences"][chain_idx] += step_stats["diverging"]
+            if not step_stats["tune"]:
+                stats["divergences"][chain_idx] += step_stats["diverging"]
+
             stats["step_size"][chain_idx] = step_stats["step_size"]
             stats["tree_size"][chain_idx] = step_stats["tree_size"]
             return stats