code: def _mask_attentions(attention, image_locs):
      batch_size, num_loc, n_att = attention.size()
      tmp1 = attention.new_zeros(num_loc)
      tmp1[:num_loc] = torch.arange(
          0, num_loc, dtype=attention.dtype
      ).unsqueeze(dim=0)
  
      tmp1 = tmp1.expand(batch_size, num_loc)
      tmp2 = image_locs.type(tmp1.type())
      tmp2 = tmp2.unsqueeze(dim=1).expand(batch_size, num_loc)
      mask = torch.ge(tmp1, tmp2)
      mask = mask.unsqueeze(dim=2).expand_as(attention)
      attention = attention.masked_fill(mask, 0)
      return attention
  
imports: import math
  from typing import Optional, Tuple, Type
  import torch
  from mmf.modules.layers import GatedTanh, ModalCombineLayer, TransformLayer
  from torch import nn
label_desc: Use image node_id 2 for calling the TopDownAttention obj's underlying member method code with example usage: obj._mask_attentions(attention, image_locs) and returns: attention