diff --git a/java/dex.go b/java/dex.go
index 983377e..2b3c931 100644
--- a/java/dex.go
+++ b/java/dex.go
@@ -252,14 +252,11 @@
 	if err != nil {
 		ctx.PropertyErrorf("min_sdk_version", "%s", err)
 	}
-	if !Bool(d.dexProperties.No_dex_container) && effectiveVersion.FinalOrFutureInt() >= 36 {
+	if !Bool(d.dexProperties.No_dex_container) && effectiveVersion.FinalOrFutureInt() >= 36 && ctx.Config().UseDexV41() {
 		// W is 36, but we have not bumped the SDK version yet, so check for both.
 		if ctx.Config().PlatformSdkVersion().FinalInt() >= 36 ||
-			ctx.Config().PlatformSdkCodename() == "Wear" {
-			// TODO(b/329465418): Skip this module since it causes issue with app DRM
-			if ctx.ModuleName() != "framework-minus-apex" {
-				flags = append([]string{"-JDcom.android.tools.r8.dexContainerExperiment"}, flags...)
-			}
+			ctx.Config().PlatformSdkCodename() == "Baklava" {
+			flags = append([]string{"-JDcom.android.tools.r8.dexContainerExperiment"}, flags...)
 		}
 	}