2 files changed, 6 insertions, 7 deletions
diff --git a/neuralnetworks/aidl/utils/src/Callbacks.cpp b/neuralnetworks/aidl/utils/src/Callbacks.cpp
index a32147734c..8084970690 100644
--- a/neuralnetworks/aidl/utils/src/Callbacks.cpp
+++ b/neuralnetworks/aidl/utils/src/Callbacks.cpp
@@ -35,7 +35,8 @@ namespace {
 
 // Converts the results of IDevice::prepareModel* to the NN canonical format. On success, this
 // function returns with a non-null nn::SharedPreparedModel with a feature level of
-// nn::Version::ANDROID_S. On failure, this function returns with the appropriate nn::GeneralError.
+// nn::kVersionFeatureLevel5. On failure, this function returns with the appropriate
+// nn::GeneralError.
 nn::GeneralResult<nn::SharedPreparedModel> prepareModelCallback(
         ErrorStatus status, const std::shared_ptr<IPreparedModel>& preparedModel) {
     HANDLE_STATUS_AIDL(status) << "model preparation failed with " << toString(status);
diff --git a/neuralnetworks/aidl/utils/src/Service.cpp b/neuralnetworks/aidl/utils/src/Service.cpp
index 01772eed53..e48593c38e 100644
--- a/neuralnetworks/aidl/utils/src/Service.cpp
+++ b/neuralnetworks/aidl/utils/src/Service.cpp
@@ -46,13 +46,11 @@ nn::GeneralResult<nn::Version> getAidlServiceFeatureLevel(IDevice* service) {
     aidlVersion = std::min(aidlVersion, IDevice::version);
 
     // Map stable AIDL versions to canonical versions.
-    switch (aidlVersion) {
-        case 1:
-            return nn::Version::ANDROID_S;
-        case 2:
-            return nn::Version::FEATURE_LEVEL_6;
+    auto version = aidlVersionToCanonicalVersion(aidlVersion);
+    if (!version.has_value()) {
+        return NN_ERROR() << "Unknown AIDL service version: " << aidlVersion;
     }
-    return NN_ERROR() << "Unknown AIDL service version: " << aidlVersion;
+    return version.value();
 }
 
 }  // namespace