diff --git a/ArmnnDriverImpl.cpp b/ArmnnDriverImpl.cpp
index 64188bb..de80445 100644
--- a/ArmnnDriverImpl.cpp
+++ b/ArmnnDriverImpl.cpp
@@ -150,7 +150,7 @@
         return ErrorStatus::NONE;
     }
 
-    unique_ptr<ArmnnPreparedModel<HalPolicy>> preparedModel(
+    sp<ArmnnPreparedModel<HalPolicy>> preparedModel(
             new ArmnnPreparedModel<HalPolicy>(
                     netId,
                     runtime.get(),
@@ -180,7 +180,7 @@
         }
     }
 
-    NotifyCallbackAndCheck(cb, ErrorStatus::NONE, preparedModel.release());
+    NotifyCallbackAndCheck(cb, ErrorStatus::NONE, preparedModel);
 
     return ErrorStatus::NONE;
 }