Bug Report for transformer-block

[raynatang]

Bug Report for https://neetcode.io/problems/transformer-block

Please describe the bug below and include any steps to reproduce the bug or screenshots if possible.

Both the official solution and my solution produce the same output that is different from the given expected output:

def init(self, model_dim, num_heads):
super().init()
torch.manual_seed(0)
self.ln1 = nn.LayerNorm(model_dim)
self.ln2 = nn.LayerNorm(model_dim)
self.multihead = self.MultiHeadedSelfAttention(model_dim, num_heads)
self.feedforward = self.VanillaNeuralNetwork(model_dim)

def forward(self, embedded):
    torch.manual_seed(0)
    first_part = self.multihead(self.ln1(embedded)) + embedded
    x = self.feedforward(self.ln2(first_part)) + first_part
    return torch.round(x, decimals=4)

Your Output:

[[[-2.1971,-2.1731,-1.1719,-0.8711],[0.6966,0.3153,-0.2754,-2.3385]],[[-0.5898,-2.4014,-0.3871,0.3969],[-0.5150,0.4576,0.9430,-0.4617]]]

Expected output:

[[[-2.3930,-2.1731,-1.3218,-0.4108],[0.6966,0.3153,-0.2754,-2.3385]],[[-0.5898,-2.4014,-0.3871,0.3969],[-0.5150,0.4576,0.9430,-0.4617]]]