diff --git a/drivers/mtd/ubi/build.c b/drivers/mtd/ubi/build.c
index 6d0d48395d8ca492215eb02239354ad21981d983..6c3fb5ab20f5723f3e97ce203bf28d274ee85f38 100644
--- a/drivers/mtd/ubi/build.c
+++ b/drivers/mtd/ubi/build.c
@@ -1019,6 +1019,8 @@ int ubi_attach_mtd_dev(struct mtd_info *mtd, int ubi_num, int vid_hdr_offset)
 out_debugfs:
 	ubi_debugfs_exit_dev(ubi);
 out_uif:
+	get_device(&ubi->dev);
+	ubi_assert(ref);
 	uif_close(ubi);
 out_detach:
 	ubi_wl_close(ubi);