diff --git a/java/dex.go b/java/dex.go
index 4a7e9dc..bc14290 100644
--- a/java/dex.go
+++ b/java/dex.go
@@ -320,7 +320,7 @@
 	// TODO(b/360905238): Remove SdkSystemServer exception after resolving missing class references.
 	if !dexParams.sdkVersion.Stable() || dexParams.sdkVersion.Kind == android.SdkSystemServer {
 		var proguardRaiseDeps classpath
-		ctx.VisitDirectDepsWithTag(proguardRaiseTag, func(m android.Module) {
+		ctx.VisitDirectDepsProxyWithTag(proguardRaiseTag, func(m android.ModuleProxy) {
 			if dep, ok := android.OtherModuleProvider(ctx, m, JavaInfoProvider); ok {
 				proguardRaiseDeps = append(proguardRaiseDeps, dep.RepackagedHeaderJars...)
 			}