fix

ml-explore · Sep 26, 2024 · 2dbda67 · 2dbda67
1 parent c76a16f
commit 2dbda67
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/mlx/backend/common/copy.cpp b/mlx/backend/common/copy.cpp
@@ -59,6 +59,12 @@ void copy_general_general(
     const std::vector<StrideT>& o_strides,
     int64_t i_offset,
     int64_t o_offset) {
+  if (data_shape.empty()) {
+    auto val = static_cast<DstT>(*(src.data<SrcT>() + i_offset));
+    auto dst_ptr = dst.data<DstT>() + o_offset;
+    *dst_ptr = val;
+    return;
+  }
   auto [shape, strides] = collapse_contiguous_dims(
       data_shape, std::vector<std::vector<StrideT>>{i_strides, o_strides});
   auto src_ptr = src.data<SrcT>() + i_offset;