Add more exceptions if the reduction fails

CM2S · May 13, 2024 · 95cf3da · 95cf3da
1 parent 9ab1b1d
commit 95cf3da
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/piglot/utils/reductions.py b/piglot/utils/reductions.py
@@ -67,7 +67,10 @@ def test_reduction(self) -> None:
                 time = torch.arange(shape[-1]).repeat(shape[:-1] + (1,))
                 data = torch.randn(*shape)
                 params = torch.randn(num_params).repeat(shape[:-1] + (1,))
-                reduced = self.reduce_torch(time, data, params)
+                try:
+                    reduced = self.reduce_torch(time, data, params)
+                except Exception as exc:
+                    raise ValueError(f"Test failed for reduction {type(self)}.") from exc
                 if reduced.shape != shape[:-1]:
                     raise ValueError(
                         f"Bad shape after reduction for {type(self)}. "