diff --git a/java/bootclasspath_fragment.go b/java/bootclasspath_fragment.go
index d6777e5..8fb8ba9 100644
--- a/java/bootclasspath_fragment.go
+++ b/java/bootclasspath_fragment.go
@@ -527,19 +527,18 @@
 	}
 
 	// Bootclasspath fragment modules that are for the platform do not produce boot related files.
-	apexInfos, _ := android.ModuleProvider(ctx, android.AllApexInfoProvider)
-	if apexInfos == nil {
+	apexInfo, _ := android.ModuleProvider(ctx, android.ApexInfoProvider)
+	if apexInfo.IsForPlatform() {
 		return ""
 	}
 
-	for _, apexInfo := range apexInfos.ApexInfos {
-		for _, apex := range apexInfo.InApexVariants {
-			if isProfileProviderApex(ctx, apex) {
-				return apex
+	for _, config := range genBootImageConfigs(ctx) {
+		if config.profileProviderModule == b.BaseModuleName() {
+			if len(config.profileImports) > 0 {
+				return config.profileImports[0]
 			}
 		}
 	}
-
 	return ""
 }