How to extract the output at the end of each TransformerDecoderLayer layer in the TransformerDecoderer

1

I am using a TransformerDecoder with 3 TransformerDecoderLayer. I was wondering if there is a way to get the output of each of the TransformerDecoderLayer?
basically:
decoder = nn.TransformerDecoder(nn.TransformerDecoderLayer(d_model=64, nhead=8), num_layers=3)
woudl it be possible to get the output of each intermediate layer (i.e. the output of first TransformerDecoderLayer, second TransformerDecoderLayer, and 3rd TransformerDecoderLayer which is equal to the output) and use it in the another model during the training?

2

You can use hooks:
https://pytorch.org/docs/master/generated/torch.nn.Module.html?highlight=hook#torch.nn.Module.register_forward_hook
(look for a good tutorial by googling)

Similarly, you ahve this new feature