can you give me your code without nest?

I tried to implement prm and countseg over 3 months. I got really tired. I implemented prm and it was very very interesting. But for countseg the prm-pytorch files are different than prm and it's pyinn that I can't generate a compatible env. Please help me or gift me the nestless code. Thank you very much for this great article