added fp16 cross entropy loss option for gpt2

22e3c7e6 · mohammad · acfe848e · 22e3c7e6
Commit 22e3c7e6 authored Jun 05, 2020 by mohammad
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

megatron/model/gpt2_model.py megatron/model/gpt2_model.py +1 -1

No files found.
--- a/megatron/model/gpt2_model.py
+++ b/megatron/model/gpt2_model.py
@@ -76,7 +76,7 @@ class GPT2Model(MegatronModule):
        if get_key_value:
            output = [output, presents]

-        if labels is not None:
+        if labels is None:
            return output
        else:
            loss = mpu.vocab_parallel_cross_entropy(output, labels)