def _manipulate_mask(self, mask: torch.BoolTensor,
student_scores: torch.Tensor,
batch: Batch) -> torch.BoolTensor:
"""
Add one extra (masked-out) token to the mask, for compatibility with BART.
"""
assert student_scores.size(1) == batch.label_vec.size(1) + 1
mask = torch.cat([mask.new_zeros([mask.size(0), 1]), mask], dim=1)
return mask
def _manipulate_mask(self, mask: torch.BoolTensor,
student_scores: torch.Tensor,
batch: Batch) -> torch.BoolTensor:
"""
Add one extra (masked-out) token to the mask, for compatibility with BART.
Only necessary when examining decoder outputs directly.
"""
if student_scores.size(1) == batch.label_vec.size(1) + 1:
mask = torch.cat([mask.new_zeros([mask.size(0), 1]), mask], dim=1)
return mask
