diff --git a/atoma-inference/src/service.rs b/atoma-inference/src/service.rs
index 33a88b34..d3e08830 100644
--- a/atoma-inference/src/service.rs
+++ b/atoma-inference/src/service.rs
@@ -213,7 +213,11 @@ mod tests {
             Ok(())
         }
 
-        fn load(_: Vec<PathBuf>, _: PrecisionBits) -> Result<Self, crate::models::ModelError> {
+        fn load(
+            _: Vec<PathBuf>,
+            _: PrecisionBits,
+            _device_id: usize,
+        ) -> Result<Self, crate::models::ModelError> {
             Ok(Self {})
         }