How to realize next stages?

[yyyyyff]

Hello, I noticed that you marked that the code is for the first stage. However, there are three stages in total in the original paper, how can I realize the next two stages by your code?? Maybe just repeat the code by three times?? Looking forward to your reply! Thank you~