follow strict tensorflow alignment requirements

wjakob · Sep 20, 2024 · ee23846 · ee23846
1 parent 4647efc
commit ee23846
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/tests/test_ndarray.cpp b/tests/test_ndarray.cpp
@@ -312,15 +312,18 @@ NB_MODULE(test_ndarray_ext, m) {
     });
 
     m.def("ret_tensorflow", []() {
-        float *f = new float[8] { 1, 2, 3, 4, 5, 6, 7, 8 };
+        struct alignas(256) Buf {
+            float f[8];
+        };
+        Buf *buf = new Buf({ 1, 2, 3, 4, 5, 6, 7, 8 });
         size_t shape[2] = { 2, 4 };
 
-        nb::capsule deleter(f, [](void *data) noexcept {
+        nb::capsule deleter(buf, [](void *data) noexcept {
            destruct_count++;
-           delete[] (float *) data;
+           delete (Buf *) data;
         });
 
-        return nb::ndarray<nb::tensorflow, float, nb::shape<2, 4>>(f, 2, shape,
+        return nb::ndarray<nb::tensorflow, float, nb::shape<2, 4>>(buf->f, 2, shape,
                                                                    deleter);
     });