Fixing bug with 0 sized tensors

coreylowman · Aug 10, 2023 · 33f7154 · 33f7154
1 parent bd6b368
commit 33f7154
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/tensor/cpu/iterate.rs b/src/tensor/cpu/iterate.rs
@@ -18,7 +18,11 @@ impl<S: Shape> NdIndex<S> {
             indices: Default::default(),
             shape: shape.concrete(),
             strides,
-            next: Some(0),
+            next: if shape.num_elements() > 0 {
+                Some(0)
+            } else {
+                None
+            },
             contiguous: (strides == shape.strides()).then(|| shape.num_elements()),
         }
     }