Commit 8f0af7cd authored by Boris Fomitchev's avatar Boris Fomitchev
Browse files

Fixed APEX test 



Signed-off-by: default avatarBoris Fomitchev <bfomitchev@nvidia.com>
parent b04eb0a2

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!')