test: zero out first input for custom-ops tests

pykeio · Dec 21, 2024 · 4d16d60 · 4d16d60
1 parent cf447cf
commit 4d16d60
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/examples/custom-ops/examples/custom-ops.rs b/examples/custom-ops/examples/custom-ops.rs
@@ -79,7 +79,13 @@ fn main() -> ort::Result<()> {
 		.commit_from_file("tests/data/custom_op_test.onnx")?;
 
 	let allocator = session.allocator();
-	let value1 = Tensor::<f32>::new(allocator, [3, 5])?;
+	let mut value1 = Tensor::<f32>::new(allocator, [3, 5])?;
+	{
+		let (_, data) = value1.extract_raw_tensor_mut();
+		for datum in data {
+			*datum = 0.;
+		}
+	}
 	let mut value2 = Tensor::<f32>::new(allocator, [3, 5])?;
 	{
 		let (_, data) = value2.extract_raw_tensor_mut();

diff --git a/src/operator/tests.rs b/src/operator/tests.rs
@@ -82,7 +82,13 @@ fn test_custom_ops() -> crate::Result<()> {
 		.commit_from_file("tests/data/custom_op_test.onnx")?;
 
 	let allocator = session.allocator();
-	let value1 = Tensor::<f32>::new(allocator, [3, 5])?;
+	let mut value1 = Tensor::<f32>::new(allocator, [3, 5])?;
+	{
+		let (_, data) = value1.extract_raw_tensor_mut();
+		for datum in data {
+			*datum = 0.;
+		}
+	}
 	let mut value2 = Tensor::<f32>::new(allocator, [3, 5])?;
 	{
 		let (_, data) = value2.extract_raw_tensor_mut();