test fixes

nv-legate · Nov 7, 2023 · d44d8ad · d44d8ad
1 parent eb71fd3
commit d44d8ad
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/src/cunumeric/matrix/batched_cholesky_template.inl b/src/cunumeric/matrix/batched_cholesky_template.inl
@@ -22,6 +22,7 @@
 #include "cunumeric/matrix/batched_cholesky.h"
 #include "cunumeric/matrix/potrf_template.inl"
 #include "cunumeric/matrix/transpose_template.inl"
+#include "cunumeric/pitches.h"
 
 namespace cunumeric {
 
@@ -67,6 +68,11 @@ struct BatchedCholeskyImpl {
         "Batched cholesky is not supported when input/output shapes differ");
     }
 
+    Pitches<DIM - 1> pitches;
+    size_t volume = pitches.flatten(shape);
+
+    if (volume == 0) return;
+
     auto ncols = shape.hi[DIM - 1] - shape.lo[DIM - 1] + 1;
 
     size_t in_strides[DIM];

diff --git a/tests/unit/cunumeric/test_config.py b/tests/unit/cunumeric/test_config.py
@@ -117,6 +117,7 @@ def test_CuNumericOpCode() -> None:
         "ADVANCED_INDEXING",
         "ARANGE",
         "ARGWHERE",
+        "BATCHED_CHOLESKY",
         "BINARY_OP",
         "BINARY_RED",
         "BINCOUNT",