How to continue on to the next task after agent reboot ##[error]We stopped hearing from agent xxx

Question

My yaml consists of 2 stages and stage A & b. one of the task under stage A has to install/uninstall the drivers and needs a system reboot. Once the self-host agent is rebooted. I want to the agent to continue to listen to next stage B. Although my agent is up after reboot and ready for listening jobs. the build that I triggered is failing with following error after sometime (approx 10-15MIN) and skipping other stages.

##[error]We stopped hearing from agent xxx. Verify the agent machine is running and has a healthy network connection. Anything that terminates an agent process, starves it for CPU, or blocks its network access can cause this error. For more information, see: https://go.microsoft.com/fwlink/?linkid=846610

Stage A:

- task: CmdLine@2
  displayName: install drivers
  inputs:
    verbose: false
    workingDirectory: $(Pipeline.Workspace)/$(Build.BuildId)/${{ parameters.targetfolder }}/A
    script: pnputil /add-driver xxx.inf /install
    failOnStdErr: false
    continueOnError: true
    
- task: CmdLine@2
  displayName: reboot
  condition: always()
  inputs:
    verbose: true
    script: |
        echo rebooting the device...
        shutdown /r /f /t 5 

**Stage B: ** . . . enter image description here

Is there a way, I can continue to next stages for the same build after agent reboots and starts listening to jobs?