code: def mask_logits_by_idx(logits, tgt_in_idx):
      # logits shape: batch_size, seq_len, candidate_size
      # tgt_in_idx shape: batch_size, seq_len
  
      # the first two symbols are reserved for padding and decoder-starting symbols
      # so they should never be a possible output label
      logits[:, :, :2] = float("-inf")
  
      batch_size, seq_len = tgt_in_idx.shape
      mask_indices = torch.tril(
          tgt_in_idx.repeat(1, seq_len).reshape(
              batch_size, seq_len, seq_len
          ),
          diagonal=0,
      )
      logits = logits.scatter(2, mask_indices, float("-inf"))
      return logits
  
imports: import copy
  import logging
  import math
  from enum import Enum
  import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 2 for calling a global function with example usage: mask_logits_by_idx(logits, tgt_in_idx) and returns: logits