Minor mcmc code cleanup

pymc-devs · Jan 9, 2025 · 3953ccc · 3953ccc
1 parent c3fb82a
commit 3953ccc
Showing 1 changed file with 1 addition and 6 deletions.
diff --git a/pymc/sampling/mcmc.py b/pymc/sampling/mcmc.py
@@ -1300,12 +1300,7 @@ def _iter_sample(
                 )
 
             yield diverging
-    except KeyboardInterrupt:
-        if isinstance(trace, ZarrChain):
-            trace.record_sampling_state(step=step)
-        trace.close()
-        raise
-    except BaseException:
+    except (KeyboardInterrupt, BaseException):
         if isinstance(trace, ZarrChain):
             trace.record_sampling_state(step=step)
         trace.close()