diff --git a/rust/image.go b/rust/image.go
index e0d267d..fec6d92 100644
--- a/rust/image.go
+++ b/rust/image.go
@@ -197,21 +197,20 @@
 	return mod.InVendor() || mod.InProduct()
 }
 
-func (mod *Module) SetImageVariation(ctx android.BaseModuleContext, variant string, module android.Module) {
-	m := module.(*Module)
+func (mod *Module) SetImageVariation(ctx android.BaseModuleContext, variant string) {
 	if variant == android.VendorRamdiskVariation {
-		m.MakeAsPlatform()
+		mod.MakeAsPlatform()
 	} else if variant == android.RecoveryVariation {
-		m.MakeAsPlatform()
+		mod.MakeAsPlatform()
 	} else if strings.HasPrefix(variant, cc.VendorVariation) {
-		m.Properties.ImageVariation = cc.VendorVariation
+		mod.Properties.ImageVariation = cc.VendorVariation
 		if strings.HasPrefix(variant, cc.VendorVariationPrefix) {
-			m.Properties.VndkVersion = strings.TrimPrefix(variant, cc.VendorVariationPrefix)
+			mod.Properties.VndkVersion = strings.TrimPrefix(variant, cc.VendorVariationPrefix)
 		}
 	} else if strings.HasPrefix(variant, cc.ProductVariation) {
-		m.Properties.ImageVariation = cc.ProductVariation
+		mod.Properties.ImageVariation = cc.ProductVariation
 		if strings.HasPrefix(variant, cc.ProductVariationPrefix) {
-			m.Properties.VndkVersion = strings.TrimPrefix(variant, cc.ProductVariationPrefix)
+			mod.Properties.VndkVersion = strings.TrimPrefix(variant, cc.ProductVariationPrefix)
 		}
 	}
 }