fixed a bug on fp16 while generating samples

cfc6924b · Mostofa Patwary · 4ad5717e · cfc6924b
Commit cfc6924b authored Feb 04, 2020 by Mostofa Patwary
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

generate_samples.py generate_samples.py +4 -0

No files found.
--- a/generate_samples.py
+++ b/generate_samples.py
@@ -98,6 +98,10 @@ def get_batch(context_tokens, args):
        args.reset_attention_mask,
        False)

+    # Fp16 conversion.
+    if args.fp16:
+        attention_mask = attention_mask.half()
+
    return tokens, attention_mask, position_ids

 def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float('Inf')):