Update

    self.sublayer = nn.ModuleList([layer for _ in range(2)])

    self.size = encoder_hsize

  def forward(

      self,

      x: torch.Tensor,

      mask: torch.Tensor,

      sa_dropout_p: float,

      adj_matrix: np.ndarray,

      distance_matrix: np.ndarray

  ):

  def forward(self, x: torch.Tensor, mask: torch.Tensor, sa_dropout_p: float,

              adj_matrix: np.ndarray, distance_matrix: np.ndarray):

    """Output computation for the MATEncoder layer.

    Parameters