text: import torch.nn as nn
  
  from .attention import MultiHeadedAttention
  from .utils import SublayerConnection, PositionwiseFeedForward
  
  
  class TransformerBlock(nn.Module):
      def __init__(self, hidden, attn_heads, feed_forward_hidden, dropout):
          super().__init__()
  
          self.attention = MultiHeadedAttention(h=attn_heads, d_model=hidden)
          self.feed_forward = PositionwiseFeedForward(
              d_model=hidden, d_ff=feed_forward_hidden, dropout=dropout)
          self.input_sublayer = SublayerConnection(size=hidden, dropout=dropout)
          self.output_sublayer = SublayerConnection(size=hidden, dropout=dropout)
          self.dropout = nn.Dropout(dropout)
  
      def forward(self, x, mask):
          x = self.input_sublayer(
              x, lambda _x: self.attention.forward(_x, _x, _x, mask=mask))
          x = self.output_sublayer(x, self.feed_forward)
          return self.dropout(x)
  