Commit 7feb02c6 authored by Jared Casper's avatar Jared Casper
Browse files

Merge branch 'apex-fix-fix' into 'master' 

Fixed APEX test

See merge request ADLR/megatron-lm!102
parents bf97f933 8f0af7cd

megatron/model/transformer.py

+1 −1
Original line number Diff line number Diff line
@@ -21,7 +21,7 @@ import torch 
try:

    from apex.normalization.fused_layer_norm import FusedLayerNorm as LayerNorm

    # Try to use FusedLayerNorm from Apex - this will trigger an error.

    _ = FusedLayerNorm(8, eps=1e-5)

    _ = LayerNorm(8, eps=1e-5)



except Exception as e:

    print('WARNING: APEX is not available, using torch.nn.LayerNorm instead of apex.normalization.FusedLayerNorm!')