diff --git a/android/config.go b/android/config.go
index dff3ea5..cb604a6 100644
--- a/android/config.go
+++ b/android/config.go
@@ -2148,6 +2148,10 @@
 	return c.productVariables.GetBuildFlagBool("RELEASE_USE_TRANSITIVE_JARS_IN_CLASSPATH")
 }
 
+func (c *config) UseDexV41() bool {
+	return c.productVariables.GetBuildFlagBool("RELEASE_USE_DEX_V41")
+}
+
 var (
 	mainlineApexContributionBuildFlagsToApexNames = map[string]string{
 		"RELEASE_APEX_CONTRIBUTIONS_ADBD":                    "com.android.adbd",
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...)
 		}
 	}