Commit 22e3c7e6 authored by mohammad's avatar mohammad
Browse files

added fp16 cross entropy loss option for gpt2

parent acfe848e

megatron/model/gpt2_model.py

+1 −1
Original line number Diff line number Diff line
@@ -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)