mlvlab / rpo Goto Github PK

  def build_attention_mask(self):
      # lazily create causal attention mask, with full attention between the vision tokens
      # pytorch uses additive attention mask; fill with -inf
      mask = torch.empty(self.context_length, self.context_length)
      mask.fill_(float("-inf"))
      mask.triu_(1)  # zero out the lower diagonal
      return mask