diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 5e71a55aec79b..2f22197219001 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -12899,11 +12899,12 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
         checkCanExecuteOrThrowUnsafe(DevicePolicyManager.OPERATION_SET_USER_RESTRICTION);
 
         if (useDevicePolicyEngine(caller, /* delegateScope= */ null)) {
+            int affectedUserId = parent ? getProfileParentId(userId) : userId;
             EnforcingAdmin admin = enforcePermissionForUserRestriction(
                     who,
                     key,
                     caller.getPackageName(),
-                    userId);
+                    affectedUserId);
             PolicyDefinition<Boolean> policyDefinition =
                     PolicyDefinition.getPolicyDefinitionForUserRestriction(key);
             if (enabledFromThisOwner) {
@@ -12911,7 +12912,7 @@ public class DevicePolicyManagerService extends IDevicePolicyManager.Stub {
                         policyDefinition,
                         admin,
                         new BooleanPolicyValue(true),
-                        parent ? getProfileParentId(userId) : userId);
+                        affectedUserId);
             } else {
                 // Remove any local and global policy that was set by the admin
                 if (!policyDefinition.isLocalOnlyPolicy()) {
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/PolicyEnforcerCallbacks.java b/services/devicepolicy/java/com/android/server/devicepolicy/PolicyEnforcerCallbacks.java
index daa8a26e00d1c..fd91249fac627 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/PolicyEnforcerCallbacks.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/PolicyEnforcerCallbacks.java
@@ -215,10 +215,9 @@ final class PolicyEnforcerCallbacks {
             }
             UserRestrictionPolicyKey parsedKey =
                     (UserRestrictionPolicyKey) policyKey;
-            // TODO: call into new UserManager API when merged
             UserManagerInternal userManager = LocalServices.getService(UserManagerInternal.class);
-//            userManager.setUserRestriction(
-//                    userId, parsedKey.getRestriction(), enabled != null && enabled);
+            userManager.setUserRestriction(
+                    userId, parsedKey.getRestriction(), enabled != null && enabled);
             return true;
         }));
     }