diff --git a/drivers/mtd/ubi/upd.c b/drivers/mtd/ubi/upd.c
index c1d7b88..9e3cd34 100644
--- a/drivers/mtd/ubi/upd.c
+++ b/drivers/mtd/ubi/upd.c
@@ -155,12 +155,12 @@ int ubi_start_update(struct ubi_device *ubi, struct ubi_volume *vol,
                if (err)
                        return err;
                vol->updating = 0;
+       } else {
+               vol->upd_buf = vmalloc(ubi->leb_size);
+               if (!vol->upd_buf)
+                       return -ENOMEM;
        }

-       vol->upd_buf = vmalloc(ubi->leb_size);
-       if (!vol->upd_buf)
-               return -ENOMEM;
-
        vol->upd_ebs = div_u64(bytes + vol->usable_leb_size - 1,
                               vol->usable_leb_size);
        vol->upd_bytes = bytes;