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