diff --git a/drivers/input/touchscreen/ad7879-spi.c b/drivers/input/touchscreen/ad7879-spi.c
index c73798297b98f9349c5151c9c2d40742c2c0ebf1..59486ccba37d7b0e6aab5a3ff48dc83ff19dbf43 100644
--- a/drivers/input/touchscreen/ad7879-spi.c
+++ b/drivers/input/touchscreen/ad7879-spi.c
@@ -49,8 +49,6 @@ static int ad7879_spi_probe(struct spi_device *spi)
 	if (IS_ERR(ts))
 		return PTR_ERR(ts);
 
-	spi_set_drvdata(spi, ts);
-
 	return 0;
 }
 
diff --git a/drivers/input/touchscreen/ad7879.c b/drivers/input/touchscreen/ad7879.c
index 1bd870277e1a18bda3c887977f8880de58160885..52daaa4edc67f5b9af3da78b459fc41e8c64df89 100644
--- a/drivers/input/touchscreen/ad7879.c
+++ b/drivers/input/touchscreen/ad7879.c
@@ -680,6 +680,8 @@ struct ad7879 *ad7879_probe(struct device *dev, struct regmap *regmap,
 	if (err)
 		return ERR_PTR(err);
 
+	dev_set_drvdata(dev, ts);
+
 	return 0;
 }
 EXPORT_SYMBOL(ad7879_probe);