Follow Figure 1 (left) for connections.

def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        """Follow Figure 1 (left) for connections."""
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)