I add a copy model to deal with p_scores forward process and it works to me in your two layers network.
I have no idea about why original code works in single fc layer, sorry.
New version (related by here):
mlp_new = deepcopy(mlp)
p_scores=mlp_new.forward(pbag[0])