Skip to content

GitLab

Commit cfc6924b authored by Mostofa Patwary's avatar Mostofa Patwary
Browse files

fixed a bug on fp16 while generating samples

parent 4ad5717e
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 0 deletions
generate_samples.py
View file @ cfc6924b
...@@ -98,6 +98,10 @@ def get_batch(context_tokens, args): ...@@ -98,6 +98,10 @@ def get_batch(context_tokens, args):
args.reset_attention_mask, args.reset_attention_mask,
False) False)
# Fp16 conversion.
if args.fp16:
attention_mask = attention_mask.half()
return tokens, attention_mask, position_ids return tokens, attention_mask, position_ids
def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')): def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):
......
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