Minor mcmc code cleanup

pymc-devs · Dec 23, 2024 · 9f4c013 · 9f4c013
1 parent 28d7bd4
commit 9f4c013
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()