code: def forward(self, query, key, value, mask=None):
      batch_size, query_len, model_dim = query.size()
      d_head = model_dim // self.heads_count
      query_projected = self.query_projection(query)
      key_projected = self.key_projection(key)
      value_projected = self.value_projection(value)
      batch_size, key_len, model_dim = key_projected.size()
      batch_size, value_len, model_dim = value_projected.size()
      query_heads = query_projected.view(
          batch_size, query_len, self.heads_count, d_head
      ).transpose(1, 2)
      key_heads = key_projected.view(
          batch_size, key_len, self.heads_count, d_head
      ).transpose(1, 2)
      value_heads = value_projected.view(
          batch_size, value_len, self.heads_count, d_head
      ).transpose(1, 2)
      attention_weights = self.scaled_dot_product(
          query_heads, key_heads
      )
      if mask is not None:
          mask_expanded = mask.unsqueeze(1).expand_as(attention_weights)
          attention_weights = attention_weights.masked_fill(
              mask_expanded, -1e18
          )
      attention = self.softmax(attention_weights)
      attention_dropped = self.dropout(attention)
      context_heads = torch.matmul(attention_dropped, value_heads)
      context_sequence = context_heads.transpose(1, 2)
      context = context_sequence.reshape(
          batch_size, query_len, model_dim
      )
      final_output = self.final_projection(context)
      return final_output
  
imports: from collections import OrderedDict
  import os
  import numpy
  import torch
  from torch import nn
  from gluoncv.torch.utils.coot_utils import truncated_normal_fill
label_desc: Use image node_id 2 for calling the MultiHeadAttention obj's underlying member method code with example usage: obj.forward(query, key, value, mask) and returns: final_output