From eebfafa58aaf77bfff78de05c4be0c0ab7c7af1b Mon Sep 17 00:00:00 2001
From: Valentin Boettcher <hiro@protagon.space>
Date: Sun, 14 Aug 2022 15:53:37 +0200
Subject: [PATCH] fix for_bath

---
 hopsflow/util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/hopsflow/util.py b/hopsflow/util.py
index 24e06ba..6b50754 100644
--- a/hopsflow/util.py
+++ b/hopsflow/util.py
@@ -127,7 +127,7 @@ class EnsembleValue:
         return len(self._value)
 
     def for_bath(self, bath: int):
-        if self.num_baths == 1 and len(self.value.shape) == 1:
+        if self.num_baths == 1 and len(self.value.shape) in [0, 1]:
             return self
 
         return EnsembleValue([(N, val[bath], σ[bath]) for N, val, σ in self._value])