WIP: fix bugs for the multiStage2

Change-Id: I1f6fea0fa7a2d485526113fdfc4a1bece8bf626c