I found that there may be ambiguity in the code, or maybe I don't understand it?

[xqy074919]

Regarding the content framed in red in the picture, I think，when block_per_stage = 3 or 4 ,block_per_stage != 1,the module of SE has become False, So the code module SE below is False no matter how block_per_stage changes.
I have a feeling that this may not be what you originally thought, I feel that by commenting out "last block" I think you may be more inclined to add attention to the last block, or to add attention a few blocks before the "last block"