code: def __init__(self, config):
      super(FMHA, self).__init__()
  
      self.p_dropout = config.attention_probs_dropout_prob
      self.h = config.num_attention_heads
      self.hidden_size = config.hidden_size
      self.d = self.hidden_size // self.h
      assert (
          self.d * self.h == self.hidden_size
      ), "Invalid hidden size/num_heads"
  
file_path: apex/apex/contrib/fmha/fmha.py
imports: import torch
  import torch.nn.functional
  import fmhalib
label_desc: Use image node_id 1 to create a new FMHA object from inherited base classes: torch.nn with example: obj = FMHA(config)