class_docstr: A slight modification of torch.nn.TransformerEncoderLayer which allows for torch.jit scripting
class_docstr_tok: ["A","slight","modification","of","torch.nn.TransformerEncoderLayer","which","allows","for","torch.jit","scripting"]
code: def forward(self, src, src_key_padding_mask=None):
      # (N, S, E) -> (S, N, E)
      src = src.transpose(0, 1)
  
      src2 = self.self_attn(
          src, src, src, key_padding_mask=src_key_padding_mask
      )[0]
      src = src + self.dropout1(src2)
      src = self.norm1(src)
      src2 = self.linear2(
          self.dropout(self.activation(self.linear1(src)))
      )
      src = src + self.dropout2(src2)
      src = self.norm2(src)
  
      # (S, N, E) -> (N, S, E)
      src = src.transpose(0, 1)
      return src
  
imports: import time
  import torch
  from torch import nn
  from hivemind.moe.server.layers.custom_experts import register_expert_class
label_desc: Use image node_id 2 for calling the TransformerEncoderLayer obj's underlying member method code with example usage: obj.forward(src, src_key_padding_mask) and returns: src