diff --git a/apex/Android.bp b/apex/Android.bp
index 21054fc..04fd5f6 100644
--- a/apex/Android.bp
+++ b/apex/Android.bp
@@ -205,11 +205,10 @@
   ],
 }
 
-filegroup {
+file_contexts {
   name: "com.android.virt-file_contexts",
-  srcs: [
-    "com.android.virt-file_contexts",
-  ],
+  defaults: ["contexts_flags_defaults"],
+  srcs: ["com.android.virt-file_contexts"],
 }
 
 filegroup {
diff --git a/apex/com.android.virt-file_contexts b/apex/com.android.virt-file_contexts
index afe9f51..78720aa 100644
--- a/apex/com.android.virt-file_contexts
+++ b/apex/com.android.virt-file_contexts
@@ -3,4 +3,6 @@
 /bin/fd_server             u:object_r:fd_server_exec:s0
 /bin/virtmgr               u:object_r:virtualizationmanager_exec:s0
 /bin/virtualizationservice u:object_r:virtualizationservice_exec:s0
-/bin/vfio_handler          u:object_r:vfio_handler_exec:s0
+is_flag_enabled(RELEASE_AVF_ENABLE_DEVICE_ASSIGNMENT, `
+    /bin/vfio_handler          u:object_r:vfio_handler_exec:s0
+')