code: def test_pruning_heads(self):
      embed_dim = 768
      num_heads = 12
      num_heads_to_keep = 8
      dummy_input = torch.randn(32, 2, embed_dim)
      mha = MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads)
      reserve_head_index = mha._get_reserve_head_index(
          num_heads_to_keep=num_heads_to_keep
      )
      mha._adaptive_prune_heads(reserve_head_index=reserve_head_index)
      mha._set_skip_embed_dim_check()
      mha(query=dummy_input, key=dummy_input, value=dummy_input)
      self.assertEqual(mha.head_dim, embed_dim / num_heads)
      self.assertEqual(mha.num_heads, num_heads_to_keep)
  
file_path: fairseq/tests/test_multihead_attention.py
imports: import random
  import unittest
  import pytest
  import torch
  from fairseq.modules.multihead_attention import MultiheadAttention, _mask_for_xformers
label_desc: Use image node_id 2 for calling the TestMultiheadAttention obj's underlying member method code with example usage: obj.test_pruning_heads() without return types