WRITE_ONCE(nd->zc_copied, true);
        }
 
-       if (refcount_dec_and_test(&uarg->refcnt)) {
-               notif->io_task_work.func = __io_notif_complete_tw;
+       if (refcount_dec_and_test(&uarg->refcnt))
                io_req_task_work_add(notif);
-       }
 }
 
 struct io_kiocb *io_alloc_notif(struct io_ring_ctx *ctx)
        notif->task = current;
        io_get_task_refs(1);
        notif->rsrc_node = NULL;
+       notif->io_task_work.func = __io_notif_complete_tw;
 
        nd = io_notif_to_data(notif);
        nd->account_pages = 0;
        struct io_notif_data *nd = io_notif_to_data(notif);
 
        /* drop slot's master ref */
-       if (refcount_dec_and_test(&nd->uarg.refcnt)) {
-               notif->io_task_work.func = __io_notif_complete_tw;
+       if (refcount_dec_and_test(&nd->uarg.refcnt))
                io_req_task_work_add(notif);
-       }
 }