Skip to content

GitLab

Commit 5637d54e authored by Myle Ott's avatar Myle Ott
Browse files

Minor fix for strip_pad functions

parent a4f86a89
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
fairseq/utils.py
View file @ 5637d54e
...@@ -236,11 +236,11 @@ def post_process_prediction(hypo_tokens, src_str, alignment, align_dict, dst_dic ...@@ -236,11 +236,11 @@ def post_process_prediction(hypo_tokens, src_str, alignment, align_dict, dst_dic
def lstrip_pad(tensor, pad): def lstrip_pad(tensor, pad):
return tensor[tensor.eq(pad).sum():] return tensor[tensor.eq(pad).long().sum():]
def rstrip_pad(tensor, pad): def rstrip_pad(tensor, pad):
strip = tensor.eq(pad).sum() strip = tensor.eq(pad).long().sum()
if strip > 0: if strip > 0:
return tensor[:-strip] return tensor[:-strip]
return tensor return tensor
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment