Is there the train model store the output in pytorch?

If you would like to store the output of a certain (internal) layer, you could use forward hooks as described here.