docstr: Args:
      input features: (B, L, I_EMBED)
  Return:
      output features: (B, H_EMBED)
docstr_tok: ["Args",":","input","features",":","(","B",",","L",",","I_EMBED",")","Return",":","output","features",":","(","B",",","H_EMBED",")"]
code: def forward(self, x, mask=None):
      """
      Args:
          input features: (B, L, I_EMBED)
      Return:
          output features: (B, H_EMBED)
      """
      # Check size of tensors
      if (
          mask == None
          or self.segment_num == None
          or self.segment_num == 1
      ):
          x = x.transpose(1, 2)  # (B, L, I_EMBED) -> (B, I_EMBED, L)
          x = self.pool(x).squeeze(
              -1
          )  # (B, I_EMBED, 1) -> (B, I_EMBED)
          return x
      else:
          B, L, I_EMBED = x.size()[:2]
          mask = (
              self.mask_embedding(mask).transpose(1, 2).unsqueeze(2)
          )  # (B, L) -> (B, L, S) -> (B, S, L) -> (B, S, 1, L)
          x = x.transpose(1, 2).unsqueeze(
              1
          )  # (B, L, I_EMBED) -> (B, I_EMBED, L) -> (B, 1, I_EMBED, L)
          x = (x * mask).view(
              [-1, I_EMBED, L]
          )  # (B, S, I_EMBED, L) -> (B * S, I_EMBED, L)
          x = self.pool(x).squeeze(
              -1
          )  # (B * S, I_EMBED, 1) -> (B * S, I_EMBED)
          x = x.view([B, -1]) - self._minus  # (B, S * I_EMBED)
          return x
  
file_path: OpenNRE/opennre/module/pool/avg_pool.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  import numpy
label_desc: Use image node_id 2 for calling the AvgPool obj's underlying member method code with example usage: obj.forward(x, mask) and returns: x, x