From 671d7047ab580b6d94b6234450a47da701dcb51a Mon Sep 17 00:00:00 2001
From: Luciano Paz <luciano.paz.neuro@gmail.com>
Date: Sun, 22 Dec 2024 23:26:38 +0100
Subject: [PATCH] Minor mcmc code cleanup

---
 pymc/sampling/mcmc.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/pymc/sampling/mcmc.py b/pymc/sampling/mcmc.py
index 83a4835737..ca91325ff1 100644
--- 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()