fix: correctly handle scalar tensors in TensorRefMut::from_raw

pykeio · Dec 26, 2024 · 73eb1d8 · 73eb1d8
1 parent 9b08ef4
commit 73eb1d8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/value/impl_tensor/create.rs b/src/value/impl_tensor/create.rs
@@ -356,7 +356,7 @@ impl<'a, T: PrimitiveTensorElementType + Debug> TensorRefMut<'a, T> {
 		let shape_ptr: *const i64 = shape.as_ptr();
 		let shape_len = shape.len();
 
-		let data_len = shape.iter().product::<i64>() as usize * std::mem::size_of::<T>();
+		let data_len = calculate_tensor_size(&shape) * std::mem::size_of::<T>();
 
 		ortsys![
 			unsafe CreateTensorWithDataAsOrtValue(