fix non-unit dim set

bayesflow-org · Oct 18, 2024 · 7e84cb4 · 7e84cb4
1 parent f4af042
commit 7e84cb4
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/bayesflow/data_adapters/transforms/as_set.py b/bayesflow/data_adapters/transforms/as_set.py
@@ -8,4 +8,7 @@ def forward(self, data: np.ndarray, **kwargs) -> np.ndarray:
         return np.atleast_3d(data)
 
     def inverse(self, data: np.ndarray, **kwargs) -> np.ndarray:
-        return np.squeeze(data, axis=2)
+        if data.shape[2] == 1:
+            return np.squeeze(data, axis=2)
+
+        return data