diff --git a/java/dexpreopt.go b/java/dexpreopt.go
index f1b7178..20dbc66 100644
--- a/java/dexpreopt.go
+++ b/java/dexpreopt.go
@@ -94,7 +94,7 @@
 	}
 
 	// Don't preopt APEX variant module
-	if am, ok := ctx.Module().(android.ApexModule); ok && !am.IsForPlatform() {
+	if apexInfo := ctx.Provider(android.ApexInfoProvider).(android.ApexInfo); !apexInfo.IsForPlatform() {
 		return true
 	}