diff --git a/MdeModulePkg/Universal/Disk/DiskIoDxe/DiskIo.c b/MdeModulePkg/Universal/Disk/DiskIoDxe/DiskIo.c
index f53c3d91d0..eb3edeb53c 100644
--- a/MdeModulePkg/Universal/Disk/DiskIoDxe/DiskIo.c
+++ b/MdeModulePkg/Universal/Disk/DiskIoDxe/DiskIo.c
@@ -1111,6 +1111,7 @@ DiskIo2FlushDiskEx (
       return Status;
     }
     Task->Signature = DISK_IO2_FLUSH_TASK_SIGNATURE;
+    Task->Token     = Token;
     Status = Private->BlockIo2->FlushBlocksEx (Private->BlockIo2, &Task->BlockIo2Token);
     if (EFI_ERROR (Status)) {
       gBS->CloseEvent (Task->BlockIo2Token.Event);