diff --git a/drivers/mtd/maps/physmap_of.c b/drivers/mtd/maps/physmap_of.c
index 39d357b..32b4a5a 100644
--- a/drivers/mtd/maps/physmap_of.c
+++ b/drivers/mtd/maps/physmap_of.c
@@ -214,12 +214,12 @@ static int __devinit of_flash_probe(struct of_device *dev,
 	if (!info)
 		goto err_out;
 
+	dev_set_drvdata(&dev->dev, info);
+
 	mtd_list = kzalloc(sizeof(struct mtd_info) * count, GFP_KERNEL);
-	if (!info)
+	if (!mtd_list)
 		goto err_out;
 
-	dev_set_drvdata(&dev->dev, info);
-
 	for (i = 0; i < count; i++) {
 		err = -ENXIO;
 		if (of_address_to_resource(dp, i, &res)) {