diff --git a/dexpreopt/dexpreopt.go b/dexpreopt/dexpreopt.go
index fc1bae1..ee04dfd 100644
--- a/dexpreopt/dexpreopt.go
+++ b/dexpreopt/dexpreopt.go
@@ -222,14 +222,6 @@
 
 	invocationPath := odexPath.ReplaceExtension(ctx, "invocation")
 
-	// TODO(skvadrik): fix this to use boot image location in the module config (currently it is broken
-	// in JIT-zygote builds, because "default" boot image is hard-coded in parts of the module config).
-	bootImage := module.DexPreoptImages[archIdx]
-	var bootImageLocation string
-	if bootImage != nil {
-		bootImageLocation = PathToLocation(bootImage, arch)
-	}
-
 	// The class loader context using paths in the build
 	var classLoaderContextHost android.Paths
 
@@ -357,7 +349,7 @@
 		Flag("--runtime-arg").FlagWithList("-Xbootclasspath-locations:", module.PreoptBootClassPathDexLocations, ":").
 		Flag("${class_loader_context_arg}").
 		Flag("${stored_class_loader_context_arg}").
-		FlagWithArg("--boot-image=", bootImageLocation).Implicits(module.DexPreoptImagesDeps[archIdx].Paths()).
+		FlagWithArg("--boot-image=", strings.Join(module.DexPreoptImageLocations, ":")).Implicits(module.DexPreoptImagesDeps[archIdx].Paths()).
 		FlagWithInput("--dex-file=", module.DexPath).
 		FlagWithArg("--dex-location=", dexLocationArg).
 		FlagWithOutput("--oat-file=", odexPath).ImplicitOutput(vdexPath).
diff --git a/java/dexpreopt_bootjars.go b/java/dexpreopt_bootjars.go
index 1d363c9..8308254 100644
--- a/java/dexpreopt_bootjars.go
+++ b/java/dexpreopt_bootjars.go
@@ -605,7 +605,6 @@
 		ctx.Strict("DEXPREOPT_IMAGE_PROFILE_BUILT_INSTALLED", image.profileInstalls.String())
 		ctx.Strict("DEXPREOPT_BOOTCLASSPATH_DEX_FILES", strings.Join(image.dexPathsDeps.Strings(), " "))
 		ctx.Strict("DEXPREOPT_BOOTCLASSPATH_DEX_LOCATIONS", strings.Join(image.dexLocationsDeps, " "))
-		ctx.Strict("DEXPREOPT_IMAGE_LOCATIONS", strings.Join(image.imageLocations, ":"))
 
 		var imageNames []string
 		for _, current := range append(d.otherImages, image) {
@@ -618,15 +617,15 @@
 			sort.Slice(arches, func(i, j int) bool { return arches[i].String() < arches[j].String() })
 
 			for _, arch := range arches {
-				ctx.Strict("DEXPREOPT_IMAGE_VDEX_BUILT_INSTALLED_"+current.name+"_"+arch.String(), current.vdexInstalls[arch].String())
-				ctx.Strict("DEXPREOPT_IMAGE_"+current.name+"_"+arch.String(), current.images[arch].String())
-				ctx.Strict("DEXPREOPT_IMAGE_DEPS_"+current.name+"_"+arch.String(), strings.Join(current.imagesDeps[arch].Strings(), " "))
-				ctx.Strict("DEXPREOPT_IMAGE_BUILT_INSTALLED_"+current.name+"_"+arch.String(), current.installs[arch].String())
-				ctx.Strict("DEXPREOPT_IMAGE_UNSTRIPPED_BUILT_INSTALLED_"+current.name+"_"+arch.String(), current.unstrippedInstalls[arch].String())
-				if current.zip != nil {
-				}
+				sfx := current.name + "_" + arch.String()
+				ctx.Strict("DEXPREOPT_IMAGE_VDEX_BUILT_INSTALLED_"+sfx, current.vdexInstalls[arch].String())
+				ctx.Strict("DEXPREOPT_IMAGE_"+sfx, current.images[arch].String())
+				ctx.Strict("DEXPREOPT_IMAGE_DEPS_"+sfx, strings.Join(current.imagesDeps[arch].Strings(), " "))
+				ctx.Strict("DEXPREOPT_IMAGE_BUILT_INSTALLED_"+sfx, current.installs[arch].String())
+				ctx.Strict("DEXPREOPT_IMAGE_UNSTRIPPED_BUILT_INSTALLED_"+sfx, current.unstrippedInstalls[arch].String())
 			}
 
+			ctx.Strict("DEXPREOPT_IMAGE_LOCATIONS_"+current.name, strings.Join(current.imageLocations, ":"))
 			ctx.Strict("DEXPREOPT_IMAGE_ZIP_"+current.name, current.zip.String())
 		}
 		ctx.Strict("DEXPREOPT_IMAGE_NAMES", strings.Join(imageNames, " "))