Update model.py

model.py

@@ -73,13 +73,12 @@ class AttentionBlock(nn.Module):
         ):  # workaround for masking bug in FA. This works because Wqkv does not have bias
             # and attention scores will be also automatically zeroed.
             u = u * padding_mask[..., None]
-        u = (
-            self.inner_mha_cls(
+        w = self.inner_mha_cls(
                 self.pre_norm(u),
                 inference_params=inference_params,
-            )
-            + u
         )
+        self.filter_output = w
+        u = w + u
         if type(padding_mask) == torch.Tensor:  # guard against bias
             u = u * padding_mask[..., None]
         u = self.mlp(self.post_norm(u)) + u