code: def mask_heads(args, model, eval_dataloader):
      """This method shows how to mask head (set some heads to zero), to test the effect on the network,
      based on the head importance scores, as described in Michel et al. (http://arxiv.org/abs/1905.10650)
      """
      _, head_importance, loss = compute_heads_importance(
          args, model, eval_dataloader, compute_entropy=False
      )
      original_score = (
          1 / loss
      )  # instead of downsteam score use the LM loss
      logger.info(
          "Pruning: original score: %f, threshold: %f",
          original_score,
          original_score * args.masking_threshold,
      )
  
      new_head_mask = torch.ones_like(head_importance)
      num_to_mask = max(
          1, int(new_head_mask.numel() * args.masking_amount)
      )
  
      current_score = original_score
      while current_score >= original_score * args.masking_threshold:
          head_mask = (
              new_head_mask.clone().detach()
          )  # save current head mask
          # heads from least important to most - keep only not-masked heads
          head_importance[head_mask == 0.0] = float("Inf")
          current_heads_to_mask = head_importance.view(-1).sort()[1]
  
          if len(current_heads_to_mask) <= num_to_mask:
              print("BREAK BY num_to_mask")
              break
  
          # mask heads
          current_heads_to_mask = current_heads_to_mask[:num_to_mask]
          logger.info(
              "Heads to mask: %s", str(current_heads_to_mask.tolist())
          )
          new_head_mask = new_head_mask.view(-1)
          new_head_mask[current_heads_to_mask] = 0.0
          new_head_mask = new_head_mask.view_as(head_mask)
          new_head_mask = new_head_mask.clone().detach()
          print_2d_tensor(new_head_mask)
  
          # Compute metric and head importance again
          _, head_importance, loss = compute_heads_importance(
              args,
              model,
              eval_dataloader,
              compute_entropy=False,
              head_mask=new_head_mask,
          )
          current_score = 1 / loss
          logger.info(
              "Masking: current score: %f, remaining heads %d (%.1f percents)",
              current_score,
              new_head_mask.sum(),
              new_head_mask.sum() / new_head_mask.numel() * 100,
          )
  
      logger.info("Final head mask")
      print_2d_tensor(head_mask)
      np.save(
          os.path.join(args.output_dir, "head_mask.npy"),
          head_mask.detach().cpu().numpy(),
      )
  
      return head_mask
  
file_path: transformers/examples/research_projects/bertology/run_prune_gpt.py
imports: import argparse
  import logging
  import os
  from datetime import datetime
  import numpy
  import torch
  from torch import nn
  from torch.utils.data import DataLoader, RandomSampler, TensorDataset
  from tqdm import tqdm
  from transformers import GPT2LMHeadModel
label_desc: Use image node_id 5 for calling a global function with example usage: mask_heads(args, model, eval_dataloader) and returns: head_mask