I have a testing code as follows, and it only prints msg from hook_fn registered with Linear.
Is there a reason why hook doesn’t work with _LinearWithBias in MultiheadAttention?

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import TransformerEncoder, TransformerEncoderLayer
class SimpleVT(nn.Module):
def __init__(self):
    super().__init__()
    self.enc_layers = TransformerEncoderLayer(40, 2, 20, 0.5)
    self.encoder = TransformerEncoder(self.enc_layers, 2)
    self.decoder = nn.Linear(40, 2)
    
def forward(self, x):
    x = self.enc_layers(x)
    x = self.encoder(x)
    x = self.decoder(x)
    return x

model = SimpleVT().cuda()
print(model)
def forward_pre_hook_fn(module, i):
print(“forward_pre_hook_fn”, type(module))
print(type(model.enc_layers.self_attn.out_proj))
model.enc_layers.self_attn.out_proj.register_forward_pre_hook(forward_pre_hook_fn)
model.enc_layers.linear1.register_forward_pre_hook(forward_pre_hook_fn)
x = torch.rand([64, 102, 40]).cuda()
y_hat = torch.rand([64, 102, 2]).cuda()
y = model(x)

There is no print from the hook registered with model.enc_layers.self_attn.out_proj.
image718×783 41.3 KB