scalar -> atleast_2d (#138)

Signed-off-by: nstarman <[email protected]>
GalOrrery · Sep 7, 2023 · 5ab435f · 5ab435f
1 parent d49368e
commit 5ab435f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/stream_ml/pytorch/prior/_track.py b/src/stream_ml/pytorch/prior/_track.py
@@ -26,7 +26,9 @@
 
 def _atleast_2d(x: Array) -> Array:
     """Ensure that x is at least 2d."""
-    if x.ndim == 1:
+    if x.ndim == 0:
+        return x[None, None]
+    elif x.ndim == 1:
         return x[:, None]
     return x