From eaa2835ac655b782960ef0d2eab5b71134df0f36 Mon Sep 17 00:00:00 2001
From: Klaus Ma <klausm@nvidia.com>
Date: Sun, 9 Feb 2025 10:03:27 +0200
Subject: [PATCH] fix err

Signed-off-by: Klaus Ma <klausm@nvidia.com>
---
 examples/pi/src/service.rs | 4 ++--
 examples/pi/src/util.rs    | 7 ++++---
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/examples/pi/src/service.rs b/examples/pi/src/service.rs
index da6ce5d..f2a49d7 100644
--- a/examples/pi/src/service.rs
+++ b/examples/pi/src/service.rs
@@ -17,7 +17,7 @@ use rand::distr::{Distribution, Uniform};
 
 use flame_rs::{
     self as flame,
-    apis::{FlameError, TaskOutput},
+    apis::{FlameError, TaskOutput, TaskInput},
     service::{SessionContext, TaskContext},
 };
 
@@ -34,7 +34,7 @@ impl flame::service::FlameService for PiService {
         let mut rng = rand::rng();
         let die = Uniform::try_from(0.0..1.0).unwrap();
 
-        let input = ctx.input.unwrap_or(util::zero_u32());
+        let input = ctx.input.unwrap_or(TaskInput::from(util::zero_u32()));
         let total = util::bytes_to_u32(input.to_vec())?;
         let mut sum = 0u32;
 
diff --git a/examples/pi/src/util.rs b/examples/pi/src/util.rs
index 4689d0f..dec950d 100644
--- a/examples/pi/src/util.rs
+++ b/examples/pi/src/util.rs
@@ -11,10 +11,11 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use flame_rs::apis::{FlameError, TaskOutput};
+use flame_rs::apis::FlameError;
 
-pub fn zero_u32() -> TaskOutput {
-    TaskOutput::from(vec![0u8, 0u8, 0u8, 0u8])
+#[allow(dead_code)]
+pub fn zero_u32() -> Vec<u8> {
+    vec![0u8, 0u8, 0u8, 0u8]
 }
 
 pub fn u32_to_bytes(i: u32) -> Vec<u8> {